Normalize Skypilot managed job ids to int (#481)

nemo_run/core/execution/skypilot_jobs.py

@@ -44,6 +44,20 @@
 logger = logging.getLogger(__name__)
 
 
+def _as_job_id(job_id: Any) -> int:
+    """Coerce a Skypilot managed job id to the int this module stores in app ids.
+
+    ``sky.jobs.client.sdk.launch`` returns ``Optional[List[int]]``, so the raw value is
+    a list such as ``[1]``. App ids written before this was handled embed that list
+    form, e.g. ``cluster___task___[1]``, so parsing accepts it too.
+    """
+    if isinstance(job_id, (list, tuple)):
+        assert len(job_id) == 1, f"Expected a single Skypilot job id, got {job_id}."
+        job_id = job_id[0]
+
+    return int(str(job_id).strip("[]"))
+
+
 @dataclass(kw_only=True)
 class SkypilotJobsExecutor(Executor):
     """
@@ -148,7 +162,7 @@ def parse_app(cls: Type["SkypilotJobsExecutor"], app_id: str) -> tuple[str, str,
         app = app_id.split("___")
         cluster, task, job_id = app[0], app[1], app[2]
         assert cluster and task and job_id, f"Invalid app id for Skypilot: {app_id}"
-        return cluster, task, int(job_id)
+        return cluster, task, _as_job_id(job_id)
 
     def to_resources(self) -> Union[set["sky.Resources"], set["sky.Resources"]]:
         from sky.resources import Resources
@@ -416,9 +430,9 @@ def launch(
         if num_nodes:
             task.num_nodes = num_nodes
 
-        job_id, handle = stream_and_get(launch(task))
+        job_ids, handle = stream_and_get(launch(task))
 
-        return job_id, handle
+        return (_as_job_id(job_ids) if job_ids is not None else None), handle
 
     def cleanup(self, handle: str):
         import sky.jobs.client.sdk as sky_jobs

test/core/execution/test_skypilot_jobs.py

@@ -150,6 +150,14 @@ def test_parse_app(self, mock_skypilot_imports):
         assert task == "task-name"
         assert job_id == 123
 
+    def test_parse_app_legacy_list_job_id(self, mock_skypilot_imports):
+        # App ids written before #481 embed Skypilot's list job id verbatim.
+        cluster, task, job_id = SkypilotJobsExecutor.parse_app("cluster-name___task-name___[123]")
+
+        assert cluster == "cluster-name"
+        assert task == "task-name"
+        assert job_id == 123
+
     def test_parse_app_invalid(self, mock_skypilot_imports):
         # The implementation raises IndexError when the app_id format is invalid
         with pytest.raises(IndexError):
@@ -324,6 +332,24 @@ def test_launch(self, mock_launch, mock_stream_and_get, executor):
         assert job_id == 123
         assert handle is mock_handle
 
+    @patch("sky.stream_and_get")
+    @patch("sky.jobs.client.sdk.launch")
+    def test_launch_normalizes_list_job_id(self, mock_launch, mock_stream_and_get, executor):
+        # sky.jobs.client.sdk.launch returns Optional[List[int]], not an int (#481).
+        mock_handle = MagicMock()
+        mock_handle.get_cluster_name.return_value = "cluster-name"
+        mock_launch.return_value = MagicMock()
+        mock_stream_and_get.return_value = ([1], mock_handle)
+
+        job_id, handle = executor.launch(MagicMock())
+
+        assert job_id == 1
+        assert handle is mock_handle
+        # The scheduler interpolates the job id straight into the app id, so it has to
+        # survive a round trip through parse_app.
+        app_id = f"{handle.get_cluster_name()}___task-name___{job_id}"
+        assert SkypilotJobsExecutor.parse_app(app_id) == ("cluster-name", "task-name", 1)
+
     def test_workdir(self, executor):
         executor.job_dir = "/path/to/job"
         assert executor.workdir == "/path/to/job/workdir"