#90 check notebook and terminal sessions

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.')
+    print("Kernel not idle. Pass.")