# Copyright (c) 2022, 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.
import logging
import threading
import time
from typing import Any, Dict, Optional
from requests import Request, RequestException, Response, Session, codes
from requests.adapters import HTTPAdapter
from nvflare.apis.overseer_spec import SP, OverseerAgent
from nvflare.security.logging import secure_format_exception
[docs]class HttpOverseerAgent(OverseerAgent):
def __init__(
self,
role,
overseer_end_point,
project,
name: str,
fl_port: str = "",
admin_port: str = "",
heartbeat_interval=5,
):
if role not in ["server", "client", "admin"]:
raise ValueError(f'Expect role in ["server", "client", "admin"] but got {role}')
super().__init__()
self._role = role
self._overseer_end_point = overseer_end_point
self._project = project
self._session = None
self._status_lock = threading.Lock()
self._report_and_query = threading.Thread(target=self._rnq_worker, args=())
self._psp = SP()
self._flag = threading.Event()
self._ca_path = None
self._cert_path = None
self._prv_key_path = None
self._last_service_session_id = ""
self._asked_to_exit = False
self._logger = logging.getLogger(self.__class__.__name__)
self._retry_delay = 4
self._asked_to_stop_retrying = False
self._update_callback = None
self._conditional_cb = False
if self._role == "server":
self._sp_end_point = ":".join([name, fl_port, admin_port])
self._heartbeat_interval = heartbeat_interval
def _send(
self, api_point, headers: Optional[Dict[str, Any]] = None, payload: Optional[Dict[str, Any]] = None
) -> Response:
try_count = 0
while not self._asked_to_stop_retrying:
try:
req = Request("POST", api_point, json=payload, headers=headers)
prepared = self._session.prepare_request(req)
resp = self._session.send(prepared)
return resp
except RequestException as e:
self._logger.debug(f"Overseer error: {secure_format_exception(e)}")
try_count += 1
time.sleep(self._retry_delay)
[docs] def set_secure_context(self, ca_path: str, cert_path: str = "", prv_key_path: str = ""):
self._ca_path = ca_path
self._cert_path = cert_path
self._prv_key_path = prv_key_path
[docs] def start(self, update_callback=None, conditional_cb=False):
self._session = Session()
adapter = HTTPAdapter(max_retries=1)
self._session.mount("http://", adapter)
self._session.mount("https://", adapter)
if self._ca_path:
self._session.verify = self._ca_path
self._session.cert = (self._cert_path, self._prv_key_path)
self._conditional_cb = conditional_cb
if update_callback:
self._update_callback = update_callback
self._report_and_query.start()
self._flag.set()
[docs] def pause(self):
self._asked_to_stop_retrying = True
self._flag.clear()
[docs] def resume(self):
self._asked_to_stop_retrying = False
self._flag.set()
[docs] def end(self):
self._asked_to_stop_retrying = True
self._flag.set()
self._asked_to_exit = True
self._report_and_query.join()
[docs] def is_shutdown(self) -> bool:
"""Return whether the agent receives a shutdown request."""
return self.overseer_info.get("system") == "shutdown"
[docs] def get_primary_sp(self) -> SP:
"""Return current primary service provider.
If primary sp not available, such as not reported by SD, connection to SD not established yet
the name and ports will be empty strings.
"""
return self._psp
[docs] def set_state(self, state) -> Response:
api_point = self._overseer_end_point + "/state"
return self._send(api_point, payload={"state": state})
def _do_callback(self):
if self._update_callback:
self._update_callback(self)
def _handle_ssid(self, ssid):
if not self._conditional_cb or self._last_service_session_id != ssid:
self._last_service_session_id = ssid
self._do_callback()
def _prepare_data(self):
data = dict(role=self._role, project=self._project)
return data
def _rnq_worker(self):
data = self._prepare_data()
if self._role == "server":
data["sp_end_point"] = self._sp_end_point
api_point = self._overseer_end_point + "/heartbeat"
while not self._asked_to_exit:
self._flag.wait()
self._rnq(api_point, headers=None, data=data)
time.sleep(self._heartbeat_interval)
def _rnq(self, api_point, headers, data):
resp = self._send(api_point, headers=headers, payload=data)
if resp is None:
return
if resp.status_code != codes.ok:
return
self.overseer_info = resp.json()
psp = self.overseer_info.get("primary_sp")
if psp:
name, fl_port, admin_port = psp.get("sp_end_point").split(":")
service_session_id = psp.get("service_session_id", "")
self._psp = SP(name, fl_port, admin_port, service_session_id, True)
# last_heartbeat = psp.get("last_heartbeat", "")
self._handle_ssid(service_session_id)
else:
self._psp = SP()
service_session_id = ""
self._handle_ssid(service_session_id)