diff --git a/scripts/auto-stop-idle/autostop.py b/scripts/auto-stop-idle/autostop.py index 643375d..6405913 100644 --- a/scripts/auto-stop-idle/autostop.py +++ b/scripts/auto-stop-idle/autostop.py @@ -11,12 +11,14 @@ # express or implied. See the License for the specific language governing # permissions and limitations under the License. -import requests +import getopt +import json +import sys from datetime import datetime -import getopt, sys -import urllib3 + import boto3 -import json +import requests +import urllib3 urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) @@ -39,10 +41,12 @@ # Read in command-line parameters idle = True -port = '8443' +port = "8443" ignore_connections = False try: - opts, args = getopt.getopt(sys.argv[1:], "ht:p:c", ["help","time=","port=","ignore-connections"]) + opts, args = getopt.getopt( + sys.argv[1:], "ht:p:c", ["help", "time=", "port=", "ignore-connections"] + ) if len(opts) == 0: raise getopt.GetoptError("No input parameters!") for opt, arg in opts: @@ -69,57 +73,79 @@ def is_idle(last_activity): - last_activity = datetime.strptime(last_activity,"%Y-%m-%dT%H:%M:%S.%fz") + last_activity = datetime.strptime(last_activity, "%Y-%m-%dT%H:%M:%S.%fz") if (datetime.now() - last_activity).total_seconds() > time: - print('Notebook is idle. Last activity time = ', last_activity) + print("Kernel is idle. Last activity time = ", last_activity) return True else: - print('Notebook is not idle. Last activity time = ', last_activity) + print("Kernel is not idle. Last activity time = ", last_activity) return False +def is_terminal_state(response): + for terminal in response: + print(terminal) + if is_idle(terminal["last_activity"]): + return True + return False + + def get_notebook_name(): - log_path = '/opt/ml/metadata/resource-metadata.json' - with open(log_path, 'r') as logs: + log_path = "/opt/ml/metadata/resource-metadata.json" + with open(log_path, "r") as logs: _logs = json.load(logs) - return _logs['ResourceName'] + return _logs["ResourceName"] + # This is hitting Jupyter's sessions API: https://github.com/jupyter/jupyter/wiki/Jupyter-Notebook-Server-API#Sessions-API -response = requests.get('https://localhost:'+port+'/api/sessions', verify=False) +response = requests.get(f"https://localhost:{port}/api/sessions", verify=False) data = response.json() if len(data) > 0: for notebook in data: # Idleness is defined by Jupyter # https://github.com/jupyter/notebook/issues/4634 - if notebook['kernel']['execution_state'] == 'idle': + if notebook["kernel"]["execution_state"] == "idle": if not ignore_connections: - if notebook['kernel']['connections'] == 0: - if not is_idle(notebook['kernel']['last_activity']): + if notebook["kernel"]["connections"] == 0: + if not is_idle(notebook["kernel"]["last_activity"]): idle = False else: idle = False - print('Notebook idle state set as %s because no kernel has been detected.' % idle) + print( + f"Notebook idle state set as {idle} because no kernel has been detected." + ) else: - if not is_idle(notebook['kernel']['last_activity']): + if not is_idle(notebook["kernel"]["last_activity"]): idle = False - print('Notebook idle state set as %s since kernel connections are ignored.' % idle) + print( + f"Notebook idle state set as {idle} since kernel connections are ignored." + ) else: - print('Notebook is not idle:', notebook['kernel']['execution_state']) + print("Notebook is not idle:", notebook["kernel"]["execution_state"]) idle = False else: - client = boto3.client('sagemaker') + client = boto3.client("sagemaker") uptime = client.describe_notebook_instance( NotebookInstanceName=get_notebook_name() - )['LastModifiedTime'] + )["LastModifiedTime"] if not is_idle(uptime.strftime("%Y-%m-%dT%H:%M:%S.%fz")): idle = False - print('Notebook idle state set as %s since no sessions detected.' % idle) + print(f"Notebook idle state set as {idle} since no sessions detected.") + +# Check terminals is idle or not +terminal_response = requests.get( + f"https://localhost:{port}/api/terminals", + verify=False, +) +terminal_data = terminal_response.json() + +terminal_idle = is_terminal_state(terminal_data) if terminal_data else True + +idle = idle and terminal_idle if idle: - print('Closing idle notebook') - client = boto3.client('sagemaker') - client.stop_notebook_instance( - NotebookInstanceName=get_notebook_name() - ) + print("Closing idle notebook") + client = boto3.client("sagemaker") + client.stop_notebook_instance(NotebookInstanceName=get_notebook_name()) else: - print('Notebook not idle. Pass.') \ No newline at end of file + print("Kernel not idle. Pass.")