From 9c7f2f310a22fb78a2ccfb8445a1cf6bd2a1bbc7 Mon Sep 17 00:00:00 2001 From: Srijan Upadhyay Date: Wed, 19 Aug 2026 13:27:07 +0530 Subject: [PATCH] fix(skypilot): normalize Skypilot managed job ids to int `sky.jobs.client.sdk.launch` returns `Optional[List[int]]`, so `SkypilotJobsExecutor.launch` handed back `[1]` while its annotation promised `Optional[int]`. The scheduler interpolates that straight into `f"{cluster}___{task}___{job_id}"`, so the app id became `cluster___task___[1]` and `parse_app` died on `int('[1]')`. The managed job kept running while nemo_run lost the handle, which also broke `nemo experiment status` and `nemo experiment logs`. Coerce the id in one place: `launch` unwraps the single-element list, and `parse_app` accepts the bracketed form so app ids already written to experiment metadata stay readable. The plain `sky.launch` path used by SkypilotExecutor returns `Optional[int]` and is untouched. Closes #481 Signed-off-by: Srijan Upadhyay --- nemo_run/core/execution/skypilot_jobs.py | 20 ++++++++++++++--- test/core/execution/test_skypilot_jobs.py | 26 +++++++++++++++++++++++ 2 files changed, 43 insertions(+), 3 deletions(-) diff --git a/nemo_run/core/execution/skypilot_jobs.py b/nemo_run/core/execution/skypilot_jobs.py index d5edb179..12077a5a 100644 --- a/nemo_run/core/execution/skypilot_jobs.py +++ b/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 diff --git a/test/core/execution/test_skypilot_jobs.py b/test/core/execution/test_skypilot_jobs.py index 66a7e953..ef2b2edc 100644 --- a/test/core/execution/test_skypilot_jobs.py +++ b/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"