def test_get_and_stop_and_kill_session(self):
session = self.ctx.get_session()
id = session.get_session_id()
session.stop()
from eggroll.core.session import ErSession
dead_session = ErSession(id)
dead_session.stop()
dead_session = ErSession(id)
dead_session.kill()
class RollSiteContext:
grpc_channel_factory = GrpcChannelFactory()
def __init__(self,
roll_site_session_id,
rp_ctx: RollPairContext,
options: dict = None):
if options is None:
options = {}
self.roll_site_session_id = roll_site_session_id
self.rp_ctx = rp_ctx
self.push_session_enabled = RollSiteConfKeys.EGGROLL_ROLLSITE_PUSH_SESSION_ENABLED.get_with(
options)
if self.push_session_enabled:
# create session for push roll_pair and object
self._push_session = ErSession(
session_id=roll_site_session_id + "_push",
options=rp_ctx.get_session().get_all_options())
self._push_rp_ctx = RollPairContext(session=self._push_session)
L.info(
f"push_session={self._push_session.get_session_id()} enabled")
def stop_push_session():
self._push_session.stop()
else:
self._push_session = None
self._push_rp_ctx = None
self.role = options["self_role"]
self.party_id = str(options["self_party_id"])
self._options = options
self._registered_comm_types = dict()
self.register_comm_type('grpc', RollSiteGrpc)
endpoint = options["proxy_endpoint"]
if isinstance(endpoint, str):
splitted = endpoint.split(':')
self.proxy_endpoint = ErEndpoint(host=splitted[0].strip(),
port=int(splitted[1].strip()))
elif isinstance(endpoint, ErEndpoint):
self.proxy_endpoint = endpoint
else:
raise ValueError("endpoint only support str and ErEndpoint type")
self.is_standalone = RollSiteConfKeys.EGGROLL_ROLLSITE_DEPLOY_MODE.get_with(
options) == "standalone"
# if self.is_standalone:
# self.stub = None
# else:
# channel = self.grpc_channel_factory.create_channel(self.proxy_endpoint)
# self.stub = proxy_pb2_grpc.DataTransferServiceStub(channel)
self.pushing_latch = CountDownLatch(0)
self.rp_ctx.get_session().add_exit_task(self._wait_push_complete)
if self.push_session_enabled:
self.rp_ctx.get_session().add_exit_task(stop_push_session)
self._wait_push_exit_timeout = int(
RollSiteConfKeys.EGGROLL_ROLLSITE_PUSH_OVERALL_TIMEOUT_SEC.
get_with(options))
L.info(f"inited RollSiteContext: {self.__dict__}")
def _wait_push_complete(self):
session_id = self.rp_ctx.get_session().get_session_id()
L.info(f"running roll site exit func for er session={session_id},"
f" roll site session id={self.roll_site_session_id}")
residual_count = self.pushing_latch.await_latch(
self._wait_push_exit_timeout)
if residual_count != 0:
L.error(
f"exit session when not finish push: "
f"residual_count={residual_count}, timeout={self._wait_push_exit_timeout}"
)
def load(self, name: str, tag: str, options: dict = None):
if options is None:
options = {}
final_options = self._options.copy()
final_options.update(options)
return RollSite(name, tag, self, options=final_options)
def register_comm_type(self, name, clazz):
self._registered_comm_types[name] = clazz
def get_comm_impl(self, name):
if name in self._registered_comm_types:
return self._registered_comm_types[name]
else:
raise ValueError(f'comm_type={name} is not registered')