Add callback_url param to session creation SDK API and CLI to suppo…

…rt session event webhook (#217)

* Add `callback_url` parameter in session creation API
* Add `callback_url` as a CLI parameter in creating a compute session

changes/217.feature

@@ -0,0 +1 @@
+Add `callback_url` parameter in session creation API/CLI to support session event webhook.

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

@@ -274,6 +274,8 @@ def prepare_mount_arg(
               help='The maximum duration to wait until the session starts.')
 @click.option('--no-reuse', is_flag=True,
               help='Do not reuse existing sessions but return an error.')
+@click.option('--callback-url', metavar='CALLBACK_URL', type=str, default=None,
+              help="Callback URL which will be called upon sesison lifecycle events.")
 # query-mode options
 @click.option('-c', '--code', metavar='CODE',
               help='The code snippet as a single string')
@@ -352,6 +354,7 @@ def prepare_mount_arg(
                    '(e.g., --assign-agent agent_id_1,agent_id_2,...)')
 def run(image, files, name,                                 # base args
         type, starts_at, enqueue_only, max_wait, no_reuse,  # job scheduling options
+        callback_url,
         code, terminal,                                     # query-mode options
         clean, build, exec, basedir,                        # batch-mode options
         env,                                                # execution environment
@@ -556,6 +559,7 @@ async def _run(session, idx, name, envs,
                 enqueue_only=enqueue_only,
                 max_wait=max_wait,
                 no_reuse=no_reuse,
+                callback_url=callback_url,
                 cluster_size=cluster_size,
                 cluster_mode=cluster_mode,
                 mounts=mount,

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

@@ -57,6 +57,8 @@ def _create_cmd(docs: str = None):
     @click.option('--depends', metavar='SESSION_ID', type=str, multiple=True,
                   help="Set the list of session ID or names that the newly created session depends on. "
                        "The session will get scheduled after all of them successfully finish.")
+    @click.option('--callback-url', metavar='CALLBACK_URL', type=str, default=None,
+                  help="Callback URL which will be called upon sesison lifecycle events.")
     # execution environment
     @click.option('-e', '--env', metavar='KEY=VAL', type=str, multiple=True,
                   help='Environment variable (may appear multiple times)')
@@ -116,6 +118,7 @@ def create(
         max_wait: bool,
         no_reuse: bool,
         depends: Sequence[str],
+        callback_url: str,
         # execution environment
         env: Sequence[str],
         # extra options
@@ -169,6 +172,7 @@ def create(
                     max_wait=max_wait,
                     no_reuse=no_reuse,
                     dependencies=depends,
+                    callback_url=callback_url,
                     cluster_size=cluster_size,
                     cluster_mode=cluster_mode,
                     mounts=mount,
@@ -258,6 +262,8 @@ def _create_from_template_cmd(docs: str = None):
     @click.option('--depends', metavar='SESSION_ID', type=str, multiple=True,
                   help="Set the list of session ID or names that the newly created session depends on. "
                        "The session will get scheduled after all of them successfully finish.")
+    @click.option('--callback-url', metavar='CALLBACK_URL', type=str, default=None,
+                  help="Callback URL which will be called upon sesison lifecycle events.")
     # execution environment
     @click.option('-e', '--env', metavar='KEY=VAL', type=str, multiple=True,
                   help='Environment variable (may appear multiple times)')
@@ -315,6 +321,7 @@ def create_from_template(
         max_wait: int | Undefined,
         no_reuse: bool,
         depends: Sequence[str],
+        callback_url: str,
         # execution environment
         env: Sequence[str],
         # extra options
@@ -370,6 +377,7 @@ def create_from_template(
                     max_wait=max_wait,
                     no_reuse=no_reuse,
                     dependencies=depends,
+                    callback_url=callback_url,
                     cluster_size=cluster_size,
                     mounts=prepared_mount,
                     mount_map=prepared_mount_map,

src/ai/backend/client/func/session.py

@@ -160,6 +160,7 @@ async def get_or_create(
         max_wait: int = 0,
         no_reuse: bool = False,
         dependencies: Sequence[str] = None,
+        callback_url: Optional[str] = None,
         mounts: List[str] = None,
         mount_map: Mapping[str, str] = None,
         envs: Mapping[str, str] = None,
@@ -275,6 +276,7 @@ async def get_or_create(
         }
         if api_session.get().api_version >= (6, '20220315'):
             params['dependencies'] = dependencies
+            params['callback_url'] = callback_url
         if api_session.get().api_version >= (6, '20200815'):
             params['clusterSize'] = cluster_size
             params['clusterMode'] = cluster_mode