From 91e8f3a8849d3ac8734b69cba046b9a7f2704b1c Mon Sep 17 00:00:00 2001
From: Kiuk Chung <kiuk@meta.com>
Date: Fri, 2 May 2025 07:24:54 -0700
Subject: [PATCH] Allow empty session name for APIs that take it (#1057)

Summary:

Makes the session name returning and accepting APIs symmetric by allowing empty session names.

Today `torchx.runner.Runner()` takes an empty string as `session_name` and returns an `app_handle` with an empty session name (e.g. `local:///app-id`) when the `run()` or `schedule()` methods are called (see the newly added unittest in `torchx/runner/test/api_test.py#test_empty_session_id()`)

However runner APIs that accept `app_handle` (e.g. `runner.cancel()`) will error when the session name is empty.  This makes the APIs non-symmetric since the `app_handle` returned by `run/schedule` APIs cannot be used.

This PR fixes this.

Reviewed By: mariusae, highker, daniel-ohayon

Differential Revision: D73799333
---
 torchx/runner/test/api_test.py | 27 +++++++++++++++++++++++++++
 torchx/specs/api.py            | 21 +++++++++++++++++++--
 torchx/specs/test/api_test.py  |  9 +++++++++
 3 files changed, 55 insertions(+), 2 deletions(-)

diff --git a/torchx/runner/test/api_test.py b/torchx/runner/test/api_test.py
index 1f9804554..49b9335ab 100644
--- a/torchx/runner/test/api_test.py
+++ b/torchx/runner/test/api_test.py
@@ -25,6 +25,7 @@
     AppDef,
     AppHandle,
     AppState,
+    parse_app_handle,
     Resource,
     Role,
     UnknownAppException,
@@ -132,6 +133,32 @@ def test_session_id(self, record_mock: MagicMock) -> None:
                 event = record_mock.call_args_list[i].args[0]
                 self.assertEqual(event.session, CURRENT_SESSION_ID)
 
+    def test_empty_session_id(self, _: MagicMock) -> None:
+        empty_session_name = ""
+        with Runner(
+            empty_session_name,
+            {"local": cast(SchedulerFactory, create_scheduler)},
+        ) as runner:
+            app = AppDef(
+                "__unused_for_test__",
+                roles=[
+                    Role(
+                        name="echo",
+                        image=str(self.tmpdir),
+                        resource=resource.SMALL,
+                        entrypoint="echo",
+                        args=["__unused_for_test__"],
+                    )
+                ],
+            )
+
+            app_handle = runner.run(app, "local", self.cfg)
+
+            scheduler, session_name, app_id = parse_app_handle(app_handle)
+            self.assertEqual(scheduler, "local")
+            self.assertEqual(empty_session_name, session_name)
+            self.assertEqual(app_handle, f"local:///{app_id}")
+
     def test_run(self, record_mock: MagicMock) -> None:
         test_file = self.tmpdir / "test_file"
 
diff --git a/torchx/specs/api.py b/torchx/specs/api.py
index 5e471d4e6..958f51943 100644
--- a/torchx/specs/api.py
+++ b/torchx/specs/api.py
@@ -1115,14 +1115,31 @@ def __init__(self, app_handle: "AppHandle") -> None:
 
 def parse_app_handle(app_handle: AppHandle) -> ParsedAppHandle:
     """
-    parses the app handle into ```(scheduler_backend, session_name, and app_id)```
+    Parses the app handle into ```(scheduler_backend, session_name, and app_id)```.
+
+    Example:
+
+    .. doctest::
+
+     assert parse_app_handle("k8s://default/foo_bar") == ("k8s", "default", "foo_bar")
+     assert parse_app_handle("k8s:///foo_bar") == ("k8s", "", "foo_bar")
+
+    Args:
+        app_handle: a URI of the form ``{scheduler}://{session_name}/{app_id}``,
+            where the ``session_name`` is optional. In this case the app handle is
+            of the form ``{scheduler}:///{app_id}`` (notice the triple slashes).
+
+    Returns: A ``Tuple`` of three elements, ``(scheduler, session_name, app_id)``
+        parsed from the app_handle URI str. If the session name is not present then
+        an empty string is returned in its place in the tuple.
+
     """
 
     # parse it manually b/c currently torchx does not
     # define allowed characters nor length for session name and app_id
     import re
 
-    pattern = r"(?P<scheduler_backend>.+)://(?P<session_name>.+)/(?P<app_id>.+)"
+    pattern = r"(?P<scheduler_backend>.+)://(?P<session_name>.*)/(?P<app_id>.+)"
     match = re.match(pattern, app_handle)
     if not match:
         raise MalformedAppHandleException(app_handle)
diff --git a/torchx/specs/test/api_test.py b/torchx/specs/test/api_test.py
index 072c7a7da..9a251c7d7 100644
--- a/torchx/specs/test/api_test.py
+++ b/torchx/specs/test/api_test.py
@@ -377,6 +377,15 @@ def test_parse_malformed_app_handles(self) -> None:
                 with self.assertRaises(MalformedAppHandleException):
                     parse_app_handle(handle)
 
+    def test_parse_app_handle_empty_session_name(self) -> None:
+        # missing session name is not OK but an empty one is
+        app_handle = "local:///my_application_id"
+        handle = parse_app_handle(app_handle)
+
+        self.assertEqual(handle.app_id, "my_application_id")
+        self.assertEqual("local", handle.scheduler_backend)
+        self.assertEqual("", handle.session_name)
+
     def test_parse(self) -> None:
         (scheduler_backend, session_name, app_id) = parse_app_handle(
             "local://my_session/my_app_id_1234"