upgrade the TTY messaging violation handler
support TTY sessions on multiple hosts
simplify the messages and support multiple violated resources
@@ -3,6 +3,7 @@ from tensorhive.core.utils.decorators import override
from typing import Generator, Dict, List, Any
from inspect import cleandoc
import logging
+from tensorhive.core import ssh
log = logging.getLogger(__name__)
@@ -10,17 +11,13 @@ class MessageSendingBehaviour():
def get_warning_message(self, data: Dict[str, Any]):
message_template = cleandoc(
'''{red_bg}{white_fg}
- You are violating {owner_name}\'s reservation!
- Please, stop all your computations immidiately.{reset}
+ You are violating the GPU reservation rules!
+ Please stop all your computations immediately.{reset}
{red_fg}{bold}
- Host: {hostname}
- Name: GPU{gpu_id}, {gpu_name}{reset}
- UUID: {gpu_uuid}
-
- Current reservation ends on {red_fg}{reservation_end}{reset}.
+ GPUs: {gpus}{reset}
If this was by a mistake, please do not do this again.
- Before starting any GPU-related computations, see TensorHive reservations calendar.
+ Before starting any GPU-related computations, check the TensorHive reservations calendar.
Regards,
TensorHive bot
@@ -28,12 +25,7 @@ class MessageSendingBehaviour():
''')
message = message_template.format(
- owner_name=data['RESERVATION_OWNER_USERNAME'],
- hostname=data['HOSTNAME'],
- gpu_id=data['GPU_ID'],
- gpu_name=data['GPU_NAME'],
- gpu_uuid=data['UUID'],
- reservation_end=data['RESERVATION_END'],
+ gpus=data['GPUS'],
white_fg=r'\e[97m',
red_fg=r'\e[31m',
light_red_fg=r'\e[101m',
@@ -50,7 +42,7 @@ class MessageSendingBehaviour():
tty=tty['TTY'])
return command
- def merged_commands(self, recipient: str, ttys: List[str], msg: str) -> str:
+ def merged_commands(self, recipient: str, ttys: List[Dict], msg: str) -> str:
'''
Concatenates multiple commands into one,
It allows using ssh connection only once instead of multiple times (for each tty separately)
@@ -65,17 +57,23 @@ class MessageSendingBehaviour():
def trigger_action(self, violation_data: Dict[str, Any]) -> None:
'''Sends a mesage to all user's terminal sessions within node'''
message = self.get_warning_message(data=violation_data)
- tty_sessions = violation_data['TTY_SESSIONS']
intruder_username = violation_data['INTRUDER_USERNAME']
- connection = violation_data['SSH_CONNECTION']
+
+ for hostname in violation_data['SSH_CONNECTIONS']:
+ node_connection = violation_data['SSH_CONNECTIONS'][hostname]
+
+ node_sessions = ssh.node_tty_sessions(node_connection)
+ tty_sessions = [sess for sess in node_sessions if sess['USER'] == intruder_username]
if not len(tty_sessions):
return
command = self.merged_commands(recipient=intruder_username, ttys=tty_sessions, msg=message)
- connection.run_command(command) # , stop_on_errors=False)
+ node_connection.run_command(command)
for tty in tty_sessions:
- log.warning('Violation warning sent to {username}, {tty_name}'.format(
+ log.warning('Violation warning sent to {username}, {tty_name}@{hostname}'.format(
username=intruder_username,
- tty_name=tty))
+ tty_name=tty,
+ hostname=hostname
+ ))
