# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Trainer-side Cell Client API for ``external_process`` and ``attach`` execution.
Rank 0 exchanges materialized tasks and results with ExternalProcessBackend. ``send()``
keeps the trainer available until all downstream result transfers settle; other ranks are
passive and rely on their training framework's collectives.
"""
import copy
import math
import os
import signal as process_signal
import threading
import time
from queue import Empty, Queue
from typing import Any, Dict, Optional
from nvflare.apis.analytix import AnalyticsDataType
from nvflare.apis.fl_constant import FLMetaKey
from nvflare.apis.shareable import Shareable
from nvflare.apis.signal import Signal
from nvflare.app_common.abstract.fl_model import FLModel, ParamsType
from nvflare.app_common.utils.fl_model_utils import FLModelUtils
from nvflare.client.api_spec import APISpec
from nvflare.client.cell.attach_session import AttachTrainerSession, TrainerSessionError
from nvflare.client.cell.bootstrap import (
ATTACH_EXECUTION_MODE,
BOOTSTRAP_FILE_ENV_VAR,
EXTERNAL_PROCESS_EXECUTION_MODE,
BootstrapKey,
get_bootstrap_client_api_type,
read_bootstrap_config,
)
from nvflare.client.cell.defs import CHANNEL, PROTOCOL_VERSION, SESSION_CONTROL_TIMEOUT, MsgKey, Topic
from nvflare.client.config import ConfigKey, ExchangeFormat, TransferType
from nvflare.client.converter_utils import convert_params
from nvflare.client.decomposers import register_framework_decomposers
from nvflare.client.utils import DIFF_FUNCS
from nvflare.fuel.f3.cellnet.cell import Cell
from nvflare.fuel.f3.cellnet.defs import MessageHeaderKey
from nvflare.fuel.f3.cellnet.defs import ReturnCode as CellReturnCode
from nvflare.fuel.f3.cellnet.fqcn import FQCN
from nvflare.fuel.f3.cellnet.utils import make_reply as make_cell_reply
from nvflare.fuel.f3.cellnet.utils import new_cell_message
from nvflare.fuel.f3.streaming.download_service import DownloadService
from nvflare.fuel.f3.streaming.shutdown import shutdown_f3_streaming
from nvflare.fuel.f3.streaming.transfer_progress import DEFAULT_STREAMING_IDLE_TIMEOUT, TransferProgressState
from nvflare.fuel.sec.authn import set_add_auth_headers_filters
from nvflare.fuel.utils.fobs import FOBSContextKey
from nvflare.fuel.utils.fobs.decomposers.via_downloader import (
RESULT_UPLOAD_PROGRESS_CTX_KEY,
RESULT_UPLOAD_TX_CREATED_CB_CTX_KEY,
ResultUploadProgressContextKey,
)
from nvflare.fuel.utils.log_utils import get_obj_logger
_HELLO_TIMEOUT = SESSION_CONTROL_TIMEOUT
_HELLO_RETRY_INTERVAL = 1.0
# Queue polling bounds abort/stop detection latency.
_RECEIVE_POLL_INTERVAL = 0.5
_HEARTBEAT_JOIN_TIMEOUT = 1.0
_OWNER_WATCHDOG_INTERVAL = 0.5
_CJ_TIMEOUT_ABORT_GRACE = 1.0
_OWNER_TERM_GRACE = 5.0
_RESULT_SOURCE_SETTLED_TIMEOUT = SESSION_CONTROL_TIMEOUT
_RESULT_SOURCE_SETTLED_ATTEMPTS = 3
_RESULT_SOURCE_SETTLED_RETRY_BACKOFF = 0.1
def _shutdown_f3_streaming() -> None:
"""Stop process-global F3 services owned by the standalone trainer."""
shutdown_f3_streaming()
def _to_python_scalar(v: Any) -> Any:
"""Convert 0-d NumPy metrics to Python scalars accepted by analytics validation."""
item = getattr(v, "item", None)
if callable(item) and getattr(v, "shape", None) == ():
return v.item()
return v
[docs]
class CellClientAPI(APISpec):
"""Client API implementation that speaks a Cell protocol to the CJ.
The control-rank task-state API is single-threaded: ``receive()``, ``send()``, and
``clear()`` must not be called concurrently.
"""
def __init__(self, bootstrap_file: Optional[str] = None):
"""Create the API from an explicit bootstrap path or the launch environment."""
super().__init__()
self.logger = get_obj_logger(self)
self._bootstrap_file = bootstrap_file or os.environ.get(BOOTSTRAP_FILE_ENV_VAR)
if not self._bootstrap_file:
raise RuntimeError(
f"no Client API Cell profile: set {BOOTSTRAP_FILE_ENV_VAR} or pass bootstrap_file "
"(external_process receives a launch bootstrap; attach uses a pre-provisioned profile)"
)
self._config = read_bootstrap_config(self._bootstrap_file)
# Legacy untyped files retain environment selection; typed envelopes must validate.
get_bootstrap_client_api_type(self._config, self._bootstrap_file)
self._execution_mode = self._config.get(
BootstrapKey.EXECUTION_MODE,
EXTERNAL_PROCESS_EXECUTION_MODE,
)
self._is_attach = self._execution_mode == ATTACH_EXECUTION_MODE
self._rank: Optional[str] = None
self._is_control_rank = False
self._cell: Optional[Cell] = None
self._session_id: Optional[str] = None
self._cj_fqcn: Optional[str] = self._config.get(BootstrapKey.CJ_FQCN)
self._cj_pid: Optional[int] = self._config.get(BootstrapKey.CJ_PID)
self._site_name: str = self._config[BootstrapKey.SITE_NAME]
self._attach = AttachTrainerSession(self) if self._is_attach else None
self._trainer_fqcn = self._attach.trainer_fqcn if self._attach else self._config[BootstrapKey.TRAINER_FQCN]
self._job_id: Optional[str] = self._config.get(BootstrapKey.JOB_ID)
# This flag controls delegated FL auth headers. Attach never receives
# the site's FL bearer credential; its separate profile SECURE_MODE
# controls Cell identity protection in AttachTrainerSession.
self._secure_mode = False if self._is_attach else bool(self._config.get(BootstrapKey.SECURE_MODE, False))
self._protocol_secure = False
self._session_security_configured = False
self._task_exchange: dict = self._config.get(BootstrapKey.TASK_EXCHANGE, {})
# Typed files predating LAUNCH_ONCE default to persistent; one-shot close is irreversible.
self._launch_once = bool(self._task_exchange.get(ConfigKey.LAUNCH_ONCE, True))
self._memory_gc_rounds = int(self._config.get(BootstrapKey.MEMORY_GC_ROUNDS, 0))
self._cuda_empty_cache = bool(self._config.get(BootstrapKey.CUDA_EMPTY_CACHE, False))
# Cell materializes FOBS payloads before the task handler queues them.
self._task_queue: "Queue[dict]" = Queue()
self._current_task: Optional[dict] = None
self._result_receiver_ids = None
self._fl_model: Optional[FLModel] = None
self._receive_called = False
self._abort = False
self._abort_reason = ""
# Task materialization may be cancelled on stop. Result publication uses a separate
# signal because SHUTDOWN can race RESULT_ACCEPTED while receivers still need the source.
self._abort_signal = Signal()
self._result_abort_signal = Signal()
self._heartbeat_cancel = Signal()
self._stopped = False
self._closed = False
self._lock = threading.Lock()
self._lifecycle_lock = threading.Lock()
self._liveness_lock = threading.Lock()
self._heartbeat_interval = 0.0
self._heartbeat_timeout = 0.0
self._last_cj_activity: Optional[float] = None
self._heartbeat_stop = threading.Event()
self._heartbeat_thread: Optional[threading.Thread] = None
self._owner_watchdog_stop = threading.Event()
self._owner_watchdog_thread: Optional[threading.Thread] = None
self._initialized = False
# An external-process trainer owns its standalone F3 runtime. Attach owns
# only this Cell session and must not tear down process-global services.
self._owns_f3_runtime = not self._is_attach
# send() owns transaction metadata; heartbeat only observes stable transfer handles.
self._result_transfer_waiters = ()
# Under _lock, this tells SHUTDOWN whether send() still owns a live result source.
self._result_send_active = False
self._params_conversion_state = {}
# ------------------------------------------------------------------ lifecycle
[docs]
def init(self, rank: Optional[str] = None):
effective_rank = rank if rank is not None else os.environ.get("RANK", "0")
with self._lifecycle_lock:
if self._closed:
self.logger.warning("called CellClientAPI.init() after shutdown; call is ignored")
return
if self._initialized:
self.logger.warning("called CellClientAPI.init() more than once; subsequent calls are ignored")
return
self._rank = effective_rank
self._is_control_rank = str(self._rank) == "0"
if not self._is_control_rank:
# Non-control ranks receive the model through framework collectives.
self._initialized = True
self.logger.info(f"rank {self._rank}: no Client API session (non-control rank)")
return
# A bare trainer must register payload decomposers normally installed by the FL process.
from nvflare.apis.utils.decomposers import flare_decomposers
from nvflare.app_common.decomposers import common_decomposers
flare_decomposers.register()
common_decomposers.register()
register_framework_decomposers(
self._task_exchange.get(ConfigKey.EXCHANGE_FORMAT, ExchangeFormat.RAW),
self._task_exchange.get(ConfigKey.SERVER_EXPECTED_FORMAT, ExchangeFormat.NUMPY),
self.logger,
)
# A failed attempt may be retried on the same API object.
self._session_id = None
self._protocol_secure = False
self._session_security_configured = False
self._heartbeat_interval = 0.0
self._heartbeat_timeout = 0.0
self._heartbeat_stop.clear()
self._heartbeat_cancel = Signal()
with self._liveness_lock:
self._last_cj_activity = None
connect_url = self._attach.prepare_connection() if self._attach else self._config[BootstrapKey.CONNECT_URL]
credentials = {}
secure = False
parent_resources = None
if self._attach:
secure, credentials = self._attach.cell_security()
parent_resources = self._attach.connection_resources()
self._protocol_secure = secure
self._cell = Cell(
fqcn=self._trainer_fqcn,
root_url=None,
# Launched mode uses the CJ's trusted local listener. Attach uses
# the independently provisioned site connection profile.
secure=secure,
credentials=credentials,
parent_url=connect_url,
parent_resources=parent_resources,
create_internal_listener=False,
auth_identity_map=self._attach.auth_identity_map() if self._attach else None,
)
try:
if self._attach:
self._attach.install_pre_decode_guard(self._cell)
# Propagate concurrent ABORT/SHUTDOWN into nested task-payload downloads.
self._cell.update_fobs_context({FOBSContextKey.ABORT_SIGNAL: self._abort_signal})
self._register_control_cbs(self._cell)
self._cell.start()
if self._attach:
self._attach.wait_for_open()
else:
self._hello()
self._start_owner_watchdog()
self._start_heartbeat()
except Exception:
self._stop_owner_watchdog()
self._stop_heartbeat()
self._stop_cell()
raise
self._initialized = True
if self._attach:
self._attach.register_cleanup()
self.logger.info(f"trainer session established: fqcn={self._trainer_fqcn} session_id={self._session_id}")
def _register_control_cbs(self, cell: Cell) -> None:
if self._attach:
self._attach.register_callbacks(cell)
cell.register_request_cb(channel=CHANNEL, topic=Topic.TASK_READY, cb=self._handle_task_ready)
cell.register_request_cb(channel=CHANNEL, topic=Topic.ABORT, cb=self._handle_abort)
cell.register_request_cb(channel=CHANNEL, topic=Topic.SHUTDOWN, cb=self._handle_shutdown)
def _hello(self) -> None:
# Cell connection is asynchronous; retry HELLO until the listener is reachable.
deadline = time.monotonic() + _HELLO_TIMEOUT
reply = None
attempt = 0
while True:
attempt += 1
remaining = deadline - time.monotonic()
if remaining <= 0:
raise TrainerSessionError(f"no HELLO reply from the CJ after {_HELLO_TIMEOUT}s (cell not connected)")
reply = self._cell.send_request(
channel=CHANNEL,
topic=Topic.HELLO,
target=self._cj_fqcn,
request=new_cell_message(
{},
{
MsgKey.TRAINER_FQCN: self._trainer_fqcn,
MsgKey.PROOF: self._config[BootstrapKey.LAUNCH_TOKEN],
# Negotiate with the trainer's compiled protocol version.
MsgKey.PROTOCOL_VERSION: PROTOCOL_VERSION,
MsgKey.JOB_ID: self._job_id,
MsgKey.SITE_NAME: self._site_name,
MsgKey.RANK: str(self._rank),
},
),
timeout=min(_HELLO_RETRY_INTERVAL, remaining),
)
rc = None if reply is None else reply.get_header(MessageHeaderKey.RETURN_CODE)
if rc == CellReturnCode.OK:
break
self.logger.debug(f"HELLO attempt {attempt} not yet delivered (rc={rc}); retrying")
time.sleep(min(_HELLO_RETRY_INTERVAL, max(0.0, deadline - time.monotonic())))
body = reply.payload
if not isinstance(body, dict) or body.get(MsgKey.REPLY_TOPIC) != Topic.HELLO_ACCEPTED:
reason = body.get(MsgKey.REASON) if isinstance(body, dict) else body
raise TrainerSessionError(f"HELLO not accepted: {reason}")
session_id = body.get(MsgKey.SESSION_ID)
if not session_id:
raise TrainerSessionError("HELLO_ACCEPTED carried no session id")
heartbeat_interval = self._valid_heartbeat_number(
MsgKey.HEARTBEAT_INTERVAL, body.get(MsgKey.HEARTBEAT_INTERVAL), positive=True
)
heartbeat_timeout = self._valid_heartbeat_number(
MsgKey.HEARTBEAT_TIMEOUT, body.get(MsgKey.HEARTBEAT_TIMEOUT), positive=False
)
if 0 < heartbeat_timeout <= heartbeat_interval:
raise TrainerSessionError(
f"invalid heartbeat policy: interval {heartbeat_interval} must be less than "
f"timeout {heartbeat_timeout}"
)
self._install_site_auth_headers(
# SECURE_MODE was added without changing protocol v1. Preserve
# compatibility with an earlier non-secure CJ that omitted it;
# a secure bootstrap still rejects False as a mismatch.
secure_mode=body.get(MsgKey.SECURE_MODE, False),
auth_token=body.get(MsgKey.AUTH_TOKEN),
token_signature=body.get(MsgKey.AUTH_TOKEN_SIGNATURE),
)
self._session_id = session_id
self._heartbeat_interval = heartbeat_interval
self._heartbeat_timeout = heartbeat_timeout
self._confirm_session_ready()
self._note_cj_activity()
def _confirm_session_ready(self) -> None:
"""Tell the CJ that HELLO_ACCEPTED processing and auth-filter setup are complete."""
started = time.monotonic()
deadline = started + _HELLO_TIMEOUT
attempt = 0
while True:
remaining = deadline - time.monotonic()
if remaining <= 0:
elapsed = time.monotonic() - started
raise TrainerSessionError(f"no SESSION_READY confirmation after {attempt} attempts over {elapsed:.1f}s")
attempt += 1
reply = self._cell.send_request(
channel=CHANNEL,
topic=Topic.SESSION_READY,
target=self._cj_fqcn,
request=new_cell_message({}, {MsgKey.SESSION_ID: self._session_id}),
timeout=min(_HELLO_RETRY_INTERVAL, remaining),
)
rc = None if reply is None else reply.get_header(MessageHeaderKey.RETURN_CODE)
body = None if reply is None else reply.payload
if (
rc == CellReturnCode.OK
and isinstance(body, dict)
and body.get(MsgKey.REPLY_TOPIC) == Topic.SESSION_READY
and body.get(MsgKey.SESSION_ID) == self._session_id
):
return
if rc == CellReturnCode.OK and isinstance(body, dict) and body.get(MsgKey.REPLY_TOPIC) == Topic.ERROR:
raise TrainerSessionError(f"SESSION_READY rejected: {body.get(MsgKey.REASON)}")
self.logger.debug(f"SESSION_READY attempt {attempt} not confirmed (rc={rc}); retrying")
time.sleep(min(_HELLO_RETRY_INTERVAL, max(0.0, deadline - time.monotonic())))
def _install_site_auth_headers(self, secure_mode, auth_token=None, token_signature=None) -> None:
if type(secure_mode) is not bool:
raise TrainerSessionError(f"session secure_mode must be a bool, got {secure_mode!r}")
if self._session_security_configured:
if secure_mode != self._secure_mode:
raise TrainerSessionError("session secure_mode changed after authentication")
return
if not self._is_attach and secure_mode != self._secure_mode:
raise TrainerSessionError("HELLO_ACCEPTED secure_mode disagrees with the launch bootstrap")
if secure_mode:
if not isinstance(auth_token, str) or not auth_token or auth_token == "NA":
raise TrainerSessionError("secure session carried no site auth token")
if not isinstance(token_signature, str) or not token_signature or token_signature == "NA":
raise TrainerSessionError("secure session carried no site auth token signature")
set_add_auth_headers_filters(
self._cell,
client_name=self._site_name,
auth_token=auth_token,
token_signature=token_signature,
)
self._secure_mode = secure_mode
self._session_security_configured = True
# ------------------------------------------------------------------ receive
[docs]
def receive(self, timeout: Optional[float] = None) -> Optional[FLModel]:
model = self._receive_internal(timeout)
if model is not None:
self._receive_called = True
return model
def _receive_internal(self, timeout: Optional[float] = None) -> Optional[FLModel]:
if not self._is_control_rank or self._closed:
return None
if self._abort:
reason = self._abort_reason
self.shutdown()
raise TrainerSessionError(f"session aborted: {reason}")
if self._stopped:
# Shut F3 down on the user thread before interpreter teardown.
self.shutdown()
return None
if self._fl_model is not None:
return self._fl_model
try:
entry = self._await_task(timeout)
except TrainerSessionError:
self.shutdown()
raise
if entry is None:
if self._stopped and not self._closed:
self.shutdown()
return None
task = entry["task"]
self._current_task = task
if self._attach:
self._attach.mark_task_delivered(task.get(MsgKey.TASK_ID))
self._result_receiver_ids = entry.get("result_receiver_ids")
self._fl_model = entry["model"]
return self._fl_model
def _await_task(self, timeout: Optional[float]) -> Optional[dict]:
"""Wait for a task; return None on SHUTDOWN/timeout and raise on ABORT."""
deadline = None if timeout is None else time.monotonic() + timeout
while True:
if self._abort:
raise TrainerSessionError(f"session aborted: {self._abort_reason}")
if self._stopped or self._closed:
return None
wait = _RECEIVE_POLL_INTERVAL
if deadline is not None:
remaining = deadline - time.monotonic()
if remaining <= 0:
return None
wait = min(wait, remaining)
try:
return self._task_queue.get(timeout=wait)
except Empty:
continue
# ------------------------------------------------------------------ send
[docs]
def send(self, model: FLModel, clear_cache: bool = True) -> None:
if not self._is_control_rank or self._closed:
return
if not self._receive_called:
raise RuntimeError('"receive" needs to be called before sending model!')
self._check_session_alive()
task = self._current_task
if task is None:
raise TrainerSessionError("send() called with no current task")
if self._task_exchange.get(ConfigKey.TRANSFER_TYPE) == TransferType.DIFF:
model = self._prepare_param_diff(model)
if model.params is None and model.metrics is None:
raise RuntimeError("the model to send does not have either params or metrics")
# DIFF is computed above in the trainer-native representation. Adapt only a
# shallow wire model so clear_cache=False leaves the user's FLModel native.
wire_model = copy.copy(model)
wire_model.params = convert_params(
model.params,
self._task_exchange.get(ConfigKey.EXCHANGE_FORMAT, ExchangeFormat.RAW),
self._task_exchange.get(ConfigKey.SERVER_EXPECTED_FORMAT, ExchangeFormat.NUMPY),
self._params_conversion_state,
self.logger,
)
shareable = FLModelUtils.to_shareable(wire_model)
def _on_transaction_created(transaction):
waiter = DownloadService.get_transfer_waiter(transaction.tx_id)
self._add_result_transfer_waiter(waiter)
# The no-op callback opts into ViaDownloader transaction tracking.
def _on_result_progress(**_kwargs):
return None
fobs_ctx_props = {
FOBSContextKey.STREAM_PROGRESS_CB: _on_result_progress,
RESULT_UPLOAD_TX_CREATED_CB_CTX_KEY: _on_transaction_created,
RESULT_UPLOAD_PROGRESS_CTX_KEY: {
ResultUploadProgressContextKey.JOB_ID: self._job_id,
ResultUploadProgressContextKey.TASK_ID: task.get(MsgKey.TASK_ID),
ResultUploadProgressContextKey.STREAMING_IDLE_TIMEOUT: DEFAULT_STREAMING_IDLE_TIMEOUT,
},
}
# Attach follows the IPC boundary: its CJ is the terminal receiver and
# materializes the result. A managed external trainer remains the source
# for the ultimate server/peer receivers propagated with TASK_READY.
source_receiver_ids = (self._cj_fqcn,) if self._attach else self._result_receiver_ids
result_accepted = False
result_id = self._attach.mark_result_publishing(task.get(MsgKey.TASK_ID)) if self._attach else None
# Serialize publication with SHUTDOWN; an admitted send owns the transfer barrier.
with self._lock:
self._check_session_alive()
self._result_send_active = True
try:
self._clear_result_transfer_waiters()
if self._attach:
self._attach.publish_result(
task_id=task.get(MsgKey.TASK_ID),
result_id=result_id,
shareable=shareable,
source_receiver_ids=source_receiver_ids,
fobs_ctx_props=fobs_ctx_props,
)
else:
request = new_cell_message(
{MessageHeaderKey.PASS_THROUGH: True},
{
MsgKey.SESSION_ID: self._session_id,
MsgKey.TASK_ID: task.get(MsgKey.TASK_ID),
MsgKey.RESULT: shareable,
},
)
reply = self._cell.send_request(
channel=CHANNEL,
topic=Topic.RESULT_READY,
target=self._cj_fqcn,
request=request,
timeout=_HELLO_TIMEOUT,
abort_signal=self._result_abort_signal,
progress_wait_cb=self._has_pending_result_transfer,
num_receivers=len(source_receiver_ids) if source_receiver_ids else 1,
receiver_ids=source_receiver_ids,
fobs_ctx_props=fobs_ctx_props,
)
self._check_result_accepted(reply)
result_accepted = True
self._note_cj_activity()
self._wait_for_result_transfers(self._snapshot_result_transfer_waiters())
if self._attach:
self._attach.mark_task_complete(task.get(MsgKey.TASK_ID))
if clear_cache:
# Acceptance and all downstream transfers succeeded; submitted and
# received parameters plus task-scoped state can now be released.
model.params = None
model.optimizer_params = None
received_model = self._fl_model
self._fl_model = None
if received_model is not None:
received_model.params = None
received_model.optimizer_params = None
self._receive_called = False
self._current_task = None
if self._attach:
self._attach.clear_result()
self._result_receiver_ids = None
except BaseException:
self._delete_result_transfers(self._snapshot_result_transfer_waiters())
if result_accepted:
# RESULT_ACCEPTED is the commit point: preserve model data for
# inspection, but do not admit a duplicate submission.
self._receive_called = False
self._current_task = None
if self._attach:
self._attach.clear_result()
self._result_receiver_ids = None
raise
finally:
try:
self._clear_result_transfer_waiters()
self._maybe_cleanup_memory()
finally:
# Downstream transfer cleanup is the send barrier. Publish that transition
# before the separately acknowledged settlement request so a concurrent
# SHUTDOWN can report the source as settled even when that request is delayed.
with self._lock:
self._result_send_active = False
if result_accepted:
self._notify_result_source_settled(task.get(MsgKey.TASK_ID))
with self._lock:
should_shutdown = self._stopped or (result_accepted and not self._launch_once)
# One-shot sessions close only after acceptance and downstream settlement.
if should_shutdown:
self.shutdown()
def _notify_result_source_settled(self, task_id: str) -> None:
"""Publish the send-barrier transition before a one-shot trainer closes its Cell."""
request = new_cell_message({}, {MsgKey.SESSION_ID: self._session_id, MsgKey.TASK_ID: task_id})
for attempt in range(_RESULT_SOURCE_SETTLED_ATTEMPTS):
try:
reply = self._cell.send_request(
channel=CHANNEL,
topic=Topic.RESULT_SOURCE_SETTLED,
target=self._cj_fqcn,
request=request,
timeout=_RESULT_SOURCE_SETTLED_TIMEOUT,
optional=True,
secure=self._protocol_secure,
)
body = reply.payload if reply is not None else None
if (
reply is not None
and reply.get_header(MessageHeaderKey.RETURN_CODE) == CellReturnCode.OK
and isinstance(body, dict)
and body.get(MsgKey.REPLY_TOPIC) == Topic.RESULT_SOURCE_SETTLED
and body.get(MsgKey.TASK_ID) == task_id
):
return
except Exception as e:
# The CJ keeps the source live until it receives this explicit
# acknowledgement. A process exit without it is treated as source
# loss, even when a launcher masks a wrapped worker failure with rc=0.
self.logger.debug(f"result-source settlement notification failed: {e}")
if attempt + 1 < _RESULT_SOURCE_SETTLED_ATTEMPTS:
time.sleep(_RESULT_SOURCE_SETTLED_RETRY_BACKOFF)
self.logger.warning(
f"result-source settlement was not acknowledged after {_RESULT_SOURCE_SETTLED_ATTEMPTS} attempts"
)
def _wait_for_result_transfers(self, result_waiters) -> None:
"""Wait for strict terminal success of every result DownloadService transaction."""
for waiter in result_waiters:
transaction_id = waiter.transaction_id
while True:
outcome = waiter.wait(timeout=_RECEIVE_POLL_INTERVAL)
if outcome is not None:
break
if waiter.done():
# The timed wait may decide to return None just before the transfer
# records its outcome and sets the event. Re-read the outcome after
# observing done so a successful completion is not mistaken for a
# terminal resolution without an outcome.
outcome = waiter.outcome
if outcome is not None:
break
raise TrainerSessionError(f"result transfer {transaction_id} ended without a terminal outcome")
if self._abort:
raise TrainerSessionError(f"session aborted while serving result: {self._abort_reason}")
if self._closed:
raise TrainerSessionError("session closed while serving result")
if outcome.status != TransferProgressState.COMPLETED:
raise TrainerSessionError(
f"result transfer {transaction_id} failed: status={outcome.status} reason={outcome.reason}"
)
@staticmethod
def _delete_result_transfers(result_waiters) -> None:
for waiter in result_waiters:
try:
DownloadService.delete_transaction(waiter.transaction_id)
except Exception:
# Preserve the original failure; idle timeout remains the cleanup backstop.
pass
def _prepare_param_diff(self, model: FLModel) -> FLModel:
exchange_format = self._task_exchange.get(ConfigKey.EXCHANGE_FORMAT, ExchangeFormat.RAW)
diff_func = DIFF_FUNCS.get(exchange_format)
if diff_func is None and exchange_format == ExchangeFormat.RAW:
diff_func = DIFF_FUNCS.get(ExchangeFormat.NUMPY)
if diff_func is None:
raise RuntimeError(f"no default params diff function for {exchange_format}")
if self._fl_model is None:
raise RuntimeError("no received model")
if self._fl_model.params is not None and model.params is not None and model.params_type == ParamsType.FULL:
try:
model.params = diff_func(original=self._fl_model.params, new=model.params)
model.params_type = ParamsType.DIFF
except Exception as e:
raise RuntimeError(f"params diff function failed: {e}") from e
return model
def _check_result_accepted(self, reply) -> None:
if reply is None:
raise TrainerSessionError("no reply to RESULT_READY from the CJ")
rc = reply.get_header(MessageHeaderKey.RETURN_CODE)
if rc != CellReturnCode.OK:
raise TrainerSessionError(f"cell-level failure on RESULT_READY: {rc}")
body = reply.payload
if not isinstance(body, dict) or body.get(MsgKey.REPLY_TOPIC) != Topic.RESULT_ACCEPTED:
reason = body.get(MsgKey.REASON) if isinstance(body, dict) else body
raise TrainerSessionError(f"result was rejected by the CJ: {reason}")
# ------------------------------------------------------------------ log / info
[docs]
def log(self, key: str, value: Any, data_type: AnalyticsDataType, **kwargs):
if self._closed:
return
self._require_control_rank("log")
try:
self._cell.fire_and_forget(
channel=CHANNEL,
topic=Topic.LOG,
targets=[self._cj_fqcn],
message=new_cell_message(
{},
{
MsgKey.SESSION_ID: self._session_id,
"key": key,
"value": _to_python_scalar(value),
"data_type": data_type,
**kwargs,
},
),
optional=True,
secure=self._protocol_secure,
)
except Exception as e:
self.logger.warning(f"failed to send LOG '{key}': {e}")
[docs]
def system_info(self) -> Dict:
return {FLMetaKey.SITE_NAME: self._site_name, FLMetaKey.JOB_ID: self._job_id}
[docs]
def get_config(self) -> Dict:
# Keep the legacy shape without exposing Cell addresses or launch credentials.
return {
ConfigKey.TASK_EXCHANGE: dict(self._task_exchange),
FLMetaKey.JOB_ID: self._job_id,
FLMetaKey.SITE_NAME: self._site_name,
ConfigKey.MEMORY_GC_ROUNDS: self._memory_gc_rounds,
ConfigKey.CUDA_EMPTY_CACHE: self._cuda_empty_cache,
}
[docs]
def get_job_id(self) -> str:
return self._job_id
[docs]
def get_site_name(self) -> str:
return self._site_name
[docs]
def get_task_name(self) -> str:
self._require_control_rank("get_task_name")
task = self._current_task
if task is None:
raise RuntimeError("no current task")
return task.get(MsgKey.TASK_NAME)
[docs]
def is_running(self) -> bool:
# Loop guards swallow session-end errors; explicit receive()/send() still raise.
if not self._is_control_rank or self._closed:
return False
if self._abort or self._stopped:
self.shutdown()
return False
try:
return self._receive_internal() is not None
except TrainerSessionError:
self.shutdown()
return False
[docs]
def is_train(self) -> bool:
self._require_control_rank("is_train")
return self._current_task_name() == self._task_exchange.get(ConfigKey.TRAIN_TASK_NAME)
[docs]
def is_evaluate(self) -> bool:
self._require_control_rank("is_evaluate")
return self._current_task_name() == self._task_exchange.get(ConfigKey.EVAL_TASK_NAME)
[docs]
def is_submit_model(self) -> bool:
self._require_control_rank("is_submit_model")
return self._current_task_name() == self._task_exchange.get(ConfigKey.SUBMIT_MODEL_TASK_NAME)
[docs]
def clear(self):
self._fl_model = None
self._receive_called = False
self._current_task = None
if self._attach:
self._attach.clear_result()
self._result_receiver_ids = None
[docs]
def shutdown(self):
"""Stop this trainer session and any process-global F3 runtime it owns.
External-process mode owns its dedicated F3 runtime and shuts it down.
Attach mode stops only the session Cell because the externally managed
process may share the process-global runtime with other work.
"""
# Stop new task admission immediately, but do not tear down the Cell or
# process-global streaming while send() still owns an accepted lazy
# result source. send() observes _stopped and performs this shutdown
# after its receiver-confirmed transfer barrier settles.
with self._lock:
self._stopped = True
defer_for_result = self._result_send_active
if defer_for_result:
self._abort_signal.trigger("client api shutdown")
return
if self._attach:
# Wake init() if it is waiting for a future job before taking the
# lifecycle lock that init() holds.
self._attach.close()
with self._lifecycle_lock:
with self._lock:
close_resources = not self._closed
if close_resources:
self._closed = True
self._stopped = True
if close_resources:
self._abort_signal.trigger("client api shutdown")
self._result_abort_signal.trigger("client api shutdown")
self._stop_owner_watchdog()
self._stop_heartbeat()
if self._owns_f3_runtime:
try:
# Keep the Cell alive until DownloadService and all retry/stream
# work have drained. An admitted reliable retry can otherwise
# enter Cell after its transport executors are shut down.
# Retry partial process-global cleanup; each operation is idempotent.
_shutdown_f3_streaming()
except Exception as e:
self.logger.warning(f"failed to stop trainer streaming services: {e}")
if close_resources:
self._stop_cell()
# ------------------------------------------------------------------ control handlers
def _handle_task_ready(self, request):
payload = request.payload
if not isinstance(payload, dict):
return make_cell_reply(CellReturnCode.INVALID_REQUEST, error="TASK_READY payload must be a dict")
reject_reason = self._validate_cj_control(request, payload)
if reject_reason:
return self._reply(Topic.TASK_FAILED, **{MsgKey.REASON: reject_reason})
task_id = payload.get(MsgKey.TASK_ID)
attempt_id = payload.get(MsgKey.ATTEMPT_ID)
task_sequence = payload.get(MsgKey.TASK_SEQ)
# Any authenticated task delivery proves that the CJ is alive. Record
# activity before attach deduplication can return an idempotent reply.
self._note_cj_activity()
if self._attach:
duplicate_reply = self._attach.reserve_task(task_id, attempt_id, task_sequence)
if duplicate_reply is not None:
return duplicate_reply
with self._lock:
terminal_reason = self._session_end_reason()
if terminal_reason:
self._forget_reserved_task(task_id, attempt_id)
return self._reply(
Topic.TASK_FAILED,
**{MsgKey.TASK_ID: task_id, MsgKey.REASON: terminal_reason},
)
shareable = payload.get(MsgKey.MODEL)
if not isinstance(shareable, Shareable):
self._forget_reserved_task(task_id, attempt_id)
return self._reply(
Topic.TASK_FAILED,
**{
MsgKey.TASK_ID: task_id,
MsgKey.REASON: f"TASK_READY model must be Shareable, got {type(shareable)}",
},
)
try:
model = FLModelUtils.from_shareable(shareable)
model.params = convert_params(
model.params,
self._task_exchange.get(ConfigKey.SERVER_EXPECTED_FORMAT, ExchangeFormat.NUMPY),
self._task_exchange.get(ConfigKey.EXCHANGE_FORMAT, ExchangeFormat.RAW),
self._params_conversion_state,
self.logger,
)
except Exception as e:
self._forget_reserved_task(task_id, attempt_id)
return self._reply(
Topic.TASK_FAILED,
**{MsgKey.TASK_ID: task_id, MsgKey.REASON: f"invalid task model: {e}"},
)
result_receiver_ids = self._normalize_result_receiver_ids(shareable.get_header(FOBSContextKey.RECEIVER_IDS))
with self._lock:
terminal_reason = self._session_end_reason()
if terminal_reason is None:
self._task_queue.put({"task": payload, "model": model, "result_receiver_ids": result_receiver_ids})
if self._attach:
self._attach.commit_reserved_task_locked(task_id, attempt_id)
if terminal_reason:
self._forget_reserved_task(task_id, attempt_id)
return self._reply(
Topic.TASK_FAILED,
**{MsgKey.TASK_ID: task_id, MsgKey.REASON: terminal_reason},
)
return self._reply(Topic.TASK_ACCEPTED, **{MsgKey.TASK_ID: task_id})
def _forget_reserved_task(self, task_id, attempt_id) -> None:
if self._attach:
self._attach.forget_reserved_task(task_id, attempt_id)
def _handle_abort(self, request):
payload = request.payload if isinstance(request.payload, dict) else {}
reject_reason = self._validate_cj_control(request, payload)
if reject_reason:
return make_cell_reply(CellReturnCode.INVALID_REQUEST, error=reject_reason)
self._note_cj_activity()
with self._lock:
self._abort = True
self._abort_reason = str(payload.get(MsgKey.REASON))
self._abort_signal.trigger(self._abort_reason)
self._result_abort_signal.trigger(self._abort_reason)
self.logger.error(f"session aborted by CJ: {self._abort_reason}")
return make_cell_reply(CellReturnCode.OK)
def _handle_shutdown(self, request):
payload = request.payload if isinstance(request.payload, dict) else {}
reject_reason = self._validate_cj_control(request, payload)
if reject_reason:
return make_cell_reply(CellReturnCode.INVALID_REQUEST, error=reject_reason)
self._note_cj_activity()
with self._lock:
self._stopped = True
result_source_live = self._result_send_active
# Cancel incoming materialization only. SHUTDOWN may race RESULT_ACCEPTED, so the
# result signal remains live until downstream transfer settlement.
self._abort_signal.trigger("session shutdown")
self.logger.info("session shutdown requested by CJ")
return make_cell_reply(CellReturnCode.OK, body={MsgKey.RESULT_SOURCE_LIVE: result_source_live})
# ------------------------------------------------------------------ helpers
@staticmethod
def _reply(reply_topic: str, **fields):
body = {MsgKey.REPLY_TOPIC: reply_topic}
body.update(fields)
return make_cell_reply(CellReturnCode.OK, body=body)
@staticmethod
def _normalize_result_receiver_ids(receiver_ids):
if isinstance(receiver_ids, str):
receiver_ids = (receiver_ids,)
if isinstance(receiver_ids, (list, tuple)):
valid = tuple(dict.fromkeys(r for r in receiver_ids if isinstance(r, str) and not FQCN.validate(r)))
if valid:
return valid
return None
def _check_session_alive(self) -> None:
reason = self._session_end_reason()
if reason:
raise TrainerSessionError(reason)
def _session_end_reason(self) -> Optional[str]:
if self._abort:
return f"session aborted: {self._abort_reason}"
if self._stopped or self._closed:
return "session stopped"
return None
def _result_publication_end_reason(self) -> Optional[str]:
"""Return only terminal conditions that may revoke an admitted result send.
Routine SHUTDOWN sets ``_stopped`` to prevent new work, but the Cell stays
alive until an already-admitted result resolves its canonical attempt and
finishes serving lazy transfers.
"""
if self._abort:
return f"session aborted: {self._abort_reason}"
if self._closed:
return "session closed"
return None
def _validate_cj_control(self, request, payload: dict) -> Optional[str]:
origin = request.get_header(MessageHeaderKey.ORIGIN) or ""
if origin != self._cj_fqcn:
return f"unexpected CJ origin {origin!r}"
if not self._session_id:
return "no active trainer session"
if payload.get(MsgKey.SESSION_ID) != self._session_id:
return "stale or unknown session id"
return None
@staticmethod
def _valid_heartbeat_number(name: str, value, positive: bool) -> float:
if (
not isinstance(value, (int, float))
or isinstance(value, bool)
or not math.isfinite(value)
or (positive and value <= 0)
or (not positive and value < 0)
):
relation = "> 0" if positive else ">= 0"
raise TrainerSessionError(f"HELLO_ACCEPTED {name} must be a finite number {relation}, got {value!r}")
return float(value)
def _start_heartbeat(self) -> None:
if self._heartbeat_timeout == 0:
return
thread = threading.Thread(target=self._heartbeat_loop, name="client_api_heartbeat", daemon=True)
self._heartbeat_thread = thread
thread.start()
def _start_owner_watchdog(self) -> None:
"""Terminate a launched trainer group if its owning CJ process disappears."""
if self._is_attach or self._cj_pid is None:
return
if not self._owner_process_alive():
# A containerized trainer can connect to its CJ while the host CJ PID is
# outside the trainer's PID namespace. Since HELLO already succeeded,
# absence on this initial probe means PID liveness is not authoritative;
# retain the Cell heartbeat/session timeout as the owner-death backstop.
self.logger.warning(
f"owning CJ process {self._cj_pid} is not visible from the trainer at session start; "
"disabling PID-based owner monitoring"
)
return
self._owner_watchdog_stop.clear()
thread = threading.Thread(
target=self._owner_watchdog_loop,
name="client_api_owner_watchdog",
daemon=True,
)
self._owner_watchdog_thread = thread
thread.start()
def _owner_watchdog_loop(self) -> None:
while not self._owner_watchdog_stop.wait(_OWNER_WATCHDOG_INTERVAL):
if self._owner_process_alive():
continue
self.logger.error(f"owning CJ process {self._cj_pid} exited; terminating external trainer process group")
self._terminate_orphaned_process_group()
return
def _owner_process_alive(self) -> bool:
try:
os.kill(self._cj_pid, 0)
return True
except ProcessLookupError:
return False
except PermissionError:
# The process still exists even if a platform denies the probe.
return True
def _terminate_orphaned_process_group(self) -> None:
"""Terminate this owned process group; hard-stop it if SIGTERM is ignored."""
if os.name != "posix":
os._exit(1)
pgid = os.getpgrp()
try:
os.killpg(pgid, process_signal.SIGTERM)
except ProcessLookupError:
return
except Exception as e:
self.logger.error(f"failed to terminate orphaned trainer process group {pgid}: {e}")
os._exit(1)
time.sleep(_OWNER_TERM_GRACE)
try:
os.killpg(pgid, process_signal.SIGKILL)
except ProcessLookupError:
return
os._exit(1)
def _stop_owner_watchdog(self) -> None:
self._owner_watchdog_stop.set()
thread = self._owner_watchdog_thread
if thread is not None and thread is not threading.current_thread() and thread.is_alive():
thread.join(timeout=_HEARTBEAT_JOIN_TIMEOUT)
def _heartbeat_loop(self) -> None:
while not self._heartbeat_stop.wait(self._heartbeat_interval):
with self._lock:
session_ended = self._closed or self._abort or (self._stopped and not self._result_send_active)
if session_ended:
return
try:
with self._lock:
result_source_live = self._result_send_active
reply = self._cell.send_request(
channel=CHANNEL,
topic=Topic.HEARTBEAT,
target=self._cj_fqcn,
request=new_cell_message(
{},
{
MsgKey.SESSION_ID: self._session_id,
MsgKey.RESULT_SOURCE_LIVE: result_source_live,
},
),
timeout=min(self._heartbeat_interval, self._heartbeat_timeout),
abort_signal=self._heartbeat_cancel,
secure=self._protocol_secure,
)
if self._heartbeat_reply_valid(reply):
self._note_cj_activity()
except Exception as e:
self.logger.debug(f"heartbeat to CJ failed: {e}")
if self._abort_if_cj_timed_out():
self._terminate_after_cj_timeout()
return
def _heartbeat_reply_valid(self, reply) -> bool:
if reply is None or reply.get_header(MessageHeaderKey.RETURN_CODE) != CellReturnCode.OK:
return False
body = reply.payload
return (
isinstance(body, dict)
and body.get(MsgKey.REPLY_TOPIC) == Topic.HEARTBEAT
and body.get(MsgKey.SESSION_ID) == self._session_id
)
def _note_cj_activity(self) -> None:
with self._liveness_lock:
self._last_cj_activity = time.monotonic()
def _add_result_transfer_waiter(self, waiter) -> None:
with self._liveness_lock:
self._result_transfer_waiters = (*self._result_transfer_waiters, waiter)
def _clear_result_transfer_waiters(self) -> None:
with self._liveness_lock:
self._result_transfer_waiters = ()
def _replace_result_transfer_waiters(self, result_waiters) -> None:
with self._liveness_lock:
self._result_transfer_waiters = tuple(result_waiters)
def _snapshot_result_transfer_waiters(self):
with self._liveness_lock:
return self._result_transfer_waiters
def _has_pending_result_transfer(self) -> bool:
with self._liveness_lock:
return any(not waiter.done() for waiter in self._result_transfer_waiters)
def _abort_if_cj_timed_out(self) -> bool:
# Use the task-state lock before the liveness lock everywhere both are needed.
# This makes heartbeat expiry atomic with TASK_READY admission.
with self._lock:
with self._liveness_lock:
if self._session_end_reason():
return False
last_activity = self._last_cj_activity
silent_for = float("inf") if last_activity is None else max(0.0, time.monotonic() - last_activity)
if silent_for <= self._heartbeat_timeout or any(
not waiter.done() for waiter in self._result_transfer_waiters
):
return False
reason = f"CJ heartbeat timed out after {silent_for:.1f}s (timeout={self._heartbeat_timeout}s)"
self._abort = True
self._abort_reason = reason
self._abort_signal.trigger(reason)
self._result_abort_signal.trigger(reason)
self._heartbeat_stop.set()
self.logger.error(reason)
return True
def _terminate_after_cj_timeout(self) -> None:
"""Escalate a lost launched session after cooperative abort has had a chance to finish."""
if self._is_attach:
return
if self._owner_watchdog_stop.wait(_CJ_TIMEOUT_ABORT_GRACE):
return
self.logger.error("CJ heartbeat timeout did not stop the trainer; terminating its process group")
self._terminate_orphaned_process_group()
def _stop_heartbeat(self) -> None:
self._heartbeat_stop.set()
self._heartbeat_cancel.trigger("client api heartbeat stopped")
thread = self._heartbeat_thread
if thread is not None and thread is not threading.current_thread() and thread.is_alive():
thread.join(timeout=_HEARTBEAT_JOIN_TIMEOUT)
def _current_task_name(self) -> Optional[str]:
task = self._current_task
return task.get(MsgKey.TASK_NAME) if task else None
def _require_control_rank(self, what: str) -> None:
if str(self._rank) != "0":
raise RuntimeError(f"only rank 0 can call {what}!")
def _stop_cell(self) -> None:
cell = self._cell
if cell is not None:
try:
cell.stop()
except Exception as e:
self.logger.debug(f"failed to stop trainer cell: {e}")