feat: Add support for SCHEDULED status (#161)

Author: achimnol

changes/161.feature

@@ -0,0 +1 @@
+Add support for the `SCHEDULED` status in `ps` (including `admin sessions`), `run`, and `start` commands

src/ai/backend/client/cli/admin/sessions.py

@@ -69,7 +69,7 @@ def transform_fields(item: SessionItem, *, in_row: bool = True) -> SessionItem:
 @admin.command()
 @click.option('-s', '--status', default=None,
               type=click.Choice([
-                  'PENDING',
+                  'PENDING', 'SCHEDULED',
                   'PREPARING', 'BUILDING', 'RUNNING', 'RESTARTING',
                   'RESIZING', 'SUSPENDED', 'TERMINATING',
                   'TERMINATED', 'ERROR', 'CANCELLED',
@@ -142,7 +142,8 @@ def sessions(status, access_key, name_only, dead, running, detail, plain, format
 
     no_match_name = None
     if status is None:
-        status = 'PENDING,PREPARING,PULLING,RUNNING,RESTARTING,TERMINATING,RESIZING,SUSPENDED,ERROR'
+        status = ('PENDING,SCHEDULED,PREPARING,PULLING,RUNNING,RESTARTING,TERMINATING,'
+                  'RESIZING,SUSPENDED,ERROR')
         no_match_name = 'active'
     if running:
         status = 'PREPARING,PULLING,RUNNING'
@@ -151,7 +152,8 @@ def sessions(status, access_key, name_only, dead, running, detail, plain, format
         status = 'CANCELLED,TERMINATED'
         no_match_name = 'dead'
     if status == 'ALL':
-        status = ('PENDING,PREPARING,PULLING,RUNNING,RESTARTING,TERMINATING,RESIZING,SUSPENDED,ERROR,'
+        status = ('PENDING,SCHEDULED,PREPARING,PULLING,RUNNING,RESTARTING,TERMINATING,'
+                  'RESIZING,SUSPENDED,ERROR,'
                   'CANCELLED,TERMINATED')
         no_match_name = 'in any status'
     if no_match_name is None:

src/ai/backend/client/cli/ps.py

@@ -7,7 +7,7 @@
 @main.command()
 @click.option('-s', '--status', default=None,
               type=click.Choice([
-                  'PENDING',
+                  'PENDING', 'SCHEDULED',
                   'PREPARING', 'BUILDING', 'RUNNING', 'RESTARTING',
                   'RESIZING', 'SUSPENDED', 'TERMINATING',
                   'TERMINATED', 'ERROR', 'CANCELLED',

src/ai/backend/client/cli/run.py

@@ -505,6 +505,10 @@ def _run_legacy(session, idx, name, envs,
             print_info('Session ID {0} is enqueued for scheduling.'
                        .format(name))
             return
+        elif compute_session.status == 'SCHEDULED':
+            print_info('Session ID {0} is scheduled and about to be started.'
+                       .format(name))
+            return
         elif compute_session.status == 'RUNNING':
             if compute_session.created:
                 vprint_done(
@@ -602,6 +606,10 @@ async def _run(session, idx, name, envs,
             print_info('Session ID {0} is enqueued for scheduling.'
                        .format(name))
             return
+        elif compute_session.status == 'SCHEDULED':
+            print_info('Session ID {0} is scheduled and about to be started.'
+                       .format(name))
+            return
         elif compute_session.status == 'RUNNING':
             if compute_session.created:
                 vprint_done(
@@ -888,6 +896,10 @@ def start(image, name, owner,                                 # base args
             if compute_session.status == 'PENDING':
                 print_info('Session ID {0} is enqueued for scheduling.'
                            .format(compute_session.id))
+            elif compute_session.status == 'SCHEDULED':
+                print_info('Session ID {0} is scheduled and about to be started.'
+                           .format(compute_session.id))
+                return
             elif compute_session.status == 'RUNNING':
                 if compute_session.created:
                     print_info('Session ID {0} is created and ready.'
@@ -1055,6 +1067,10 @@ def start_template(
             if compute_session.status == 'PENDING':
                 print_info('Session ID {0} is enqueued for scheduling.'
                            .format(name))
+            elif compute_session.status == 'SCHEDULED':
+                print_info('Session ID {0} is scheduled and about to be started.'
+                           .format(name))
+                return
             elif compute_session.status == 'RUNNING':
                 if compute_session.created:
                     print_info('Session ID {0} is created and ready.'