Add support for job groups for local executor (#220)

* Add support for job groups for local executor

Signed-off-by: Hemil Desai <[email protected]>

* handle proper exit

Signed-off-by: Hemil Desai <[email protected]>

* fix

Signed-off-by: Hemil Desai <[email protected]>

* fix

Signed-off-by: Hemil Desai <[email protected]>

* fix

Signed-off-by: Hemil Desai <[email protected]>

---------

Signed-off-by: Hemil Desai <[email protected]>

Author: hemildesai

nemo_run/run/job.py

@@ -9,6 +9,7 @@
 from nemo_run.config import Config, ConfigurableMixin, Partial, Script
 from nemo_run.core.execution.base import Executor
 from nemo_run.core.execution.docker import DockerExecutor
+from nemo_run.core.execution.local import LocalExecutor
 from nemo_run.core.execution.slurm import SlurmExecutor
 from nemo_run.core.frontend.console.api import CONSOLE
 from nemo_run.core.serialization.zlib_json import ZlibJSONSerializer
@@ -211,7 +212,7 @@ class JobGroup(ConfigurableMixin):
     handle groups of related tasks.
     """
 
-    SUPPORTED_EXECUTORS = [SlurmExecutor, DockerExecutor]
+    SUPPORTED_EXECUTORS = [SlurmExecutor, DockerExecutor, LocalExecutor]
 
     id: str
     tasks: list[Union[Partial, Script]]
@@ -370,16 +371,19 @@ def launch(
                 self.launched = True
 
     def wait(self, runner: Runner | None = None):
-        assert len(self.handles) == 1, "Only one handle is supported for task groups currently."
+        new_states = []
         try:
-            status = wait_and_exit(
-                app_handle=self.handles[0],
-                log=self.tail_logs,
-                runner=runner,
-            )
-            self.states = [status.state]
+            for handle in self.handles:
+                status = wait_and_exit(
+                    app_handle=handle,
+                    log=self.tail_logs,
+                    runner=runner,
+                )
+                new_states.append(status.state)
         except nemo_run.exceptions.UnknownStatusError:
-            self.states = [AppState.UNKNOWN]
+            new_states = [AppState.UNKNOWN]
+
+        self.states = new_states
 
     def cancel(self, runner: Runner):
         if not self.handles:

nemo_run/run/torchx_backend/launcher.py

@@ -122,6 +122,7 @@ def wait_and_exit(
     log: bool,
     runner: Runner | None = None,
     timeout: int = 10,
+    log_join_timeout: int = 600,
 ) -> specs.AppStatus:
     if runner is None:
         runner = get_runner()
@@ -159,6 +160,12 @@ def wait_and_exit(
 
     logger.info(f"Job {app_id} finished: {status.state}")
 
+    if log_thread and log_thread.is_alive():
+        logger.debug("Waiting for log thread to complete...")
+        log_thread.join(timeout=log_join_timeout)
+        if log_thread.is_alive():
+            logger.warning("Log thread did not complete within timeout, some logs may be missing")
+
     return status
 
 

nemo_run/run/torchx_backend/schedulers/local.py

@@ -44,11 +44,12 @@
     PopenRequest,
     _LocalAppDef,
 )
-from torchx.specs.api import AppDef, AppState, Role
+from torchx.specs.api import AppDef, AppState, Role, is_terminal, parse_app_handle
 
 from nemo_run.config import get_nemorun_home
 from nemo_run.core.execution.base import Executor
 from nemo_run.core.execution.local import LocalExecutor
+from nemo_run.run import experiment as run_experiment
 from nemo_run.run.torchx_backend.schedulers.api import SchedulerMixin
 
 try:
@@ -69,6 +70,7 @@ def __init__(
         image_provider_class: Callable[[LocalOpts], ImageProvider],
         cache_size: int = 100,
         extra_paths: Optional[list[str]] = None,
+        experiment: Optional[run_experiment.Experiment] = None,
     ) -> None:
         # NOTE: make sure any new init options are supported in create_scheduler(...)
         self.backend = "local"
@@ -86,6 +88,7 @@ def __init__(
         # sets lazily on submit or dryrun based on log_dir cfg
         self._base_log_dir: Optional[str] = None
         self._created_tmp_log_dir: bool = False
+        self.experiment = experiment
 
     def _submit_dryrun(self, app: AppDef, cfg: Executor) -> AppDryRunInfo[PopenRequest]:  # type: ignore
         assert isinstance(cfg, LocalExecutor), f"{cfg.__class__} not supported for local scheduler."
@@ -106,9 +109,31 @@ def schedule(self, dryrun_info: AppDryRunInfo[PopenRequest]) -> str:
 
     def describe(self, app_id: str) -> Optional[DescribeAppResponse]:
         resp = super().describe(app_id=app_id)
-
         if resp:
             _save_job_dir(self._apps)
+            if self.experiment:
+                maybe_job_to_kill = None
+                for job in self.experiment.jobs:
+                    if isinstance(job, run_experiment.JobGroup):
+                        for handle in job.handles:
+                            _, _, job_id = parse_app_handle(handle)
+                            if job_id == app_id:
+                                maybe_job_to_kill = job
+                                break
+
+                if maybe_job_to_kill:
+                    to_kill = False
+                    for handle in maybe_job_to_kill.handles:
+                        _, _, _id = parse_app_handle(handle)
+                        resp = super().describe(app_id=_id)
+                        if resp and is_terminal(resp.state):
+                            to_kill = True
+
+                    if to_kill:
+                        for handle in maybe_job_to_kill.handles:
+                            _, _, _id = parse_app_handle(handle)
+                            self._apps[_id].kill()
+
             return resp
 
         saved_apps = _get_job_dirs()
@@ -173,11 +198,14 @@ def create_scheduler(
     extra_paths: Optional[list[str]] = None,
     **kwargs: Any,
 ) -> PersistentLocalScheduler:
+    experiment = kwargs.pop("experiment", None)
+
     return PersistentLocalScheduler(
         session_name=session_name,
         image_provider_class=CWDImageProvider,
         cache_size=cache_size,
         extra_paths=extra_paths,
+        experiment=experiment,
     )
 
 

test/core/execution/test_dgxcloud.py

@@ -20,9 +20,9 @@
 
 import pytest
 
+from nemo_run.config import set_nemorun_home
 from nemo_run.core.execution.dgxcloud import DGXCloudExecutor, DGXCloudState
 from nemo_run.core.packaging.git import GitArchivePackager
-from nemo_run.config import set_nemorun_home
 
 
 class TestDGXCloudExecutor:
@@ -252,17 +252,20 @@ def test_delete_workload(self, mock_delete):
             headers=executor._default_headers(token="test_token"),
         )
 
+    @patch("time.sleep")
     @patch.object(DGXCloudExecutor, "create_data_mover_workload")
     @patch.object(DGXCloudExecutor, "status")
     @patch.object(DGXCloudExecutor, "delete_workload")
-    def test_move_data_success(self, mock_delete, mock_status, mock_create):
+    def test_move_data_success(self, mock_delete, mock_status, mock_create, mock_sleep):
         mock_response = MagicMock()
         mock_response.status_code = 200
         mock_response.json.return_value = {"workloadId": "job123", "actualPhase": "Pending"}
         mock_create.return_value = mock_response
         mock_delete.return_value = mock_response
 
-        mock_status.return_value = DGXCloudState.COMPLETED
+        # Set up status to change after first check to avoid infinite loop
+        # First return PENDING, then return COMPLETED
+        mock_status.side_effect = [DGXCloudState.PENDING, DGXCloudState.COMPLETED]
 
         executor = DGXCloudExecutor(
             base_url="https://dgxapi.example.com",
@@ -276,12 +279,17 @@ def test_move_data_success(self, mock_delete, mock_status, mock_create):
 
         executor.move_data(token="test_token", project_id="proj_id", cluster_id="cluster_id")
 
+        # Verify all expected calls were made
         mock_create.assert_called_once_with("test_token", "proj_id", "cluster_id")
         mock_status.assert_called()
         mock_delete.assert_called_once_with("test_token", "job123")
 
+        # Verify time.sleep was called
+        mock_sleep.assert_called()
+
+    @patch("time.sleep")
     @patch.object(DGXCloudExecutor, "create_data_mover_workload")
-    def test_move_data_data_mover_fail(self, mock_create):
+    def test_move_data_data_mover_fail(self, mock_create, mock_sleep):
         mock_response = MagicMock()
         mock_response.status_code = 400
 
@@ -298,9 +306,10 @@ def test_move_data_data_mover_fail(self, mock_create):
         with pytest.raises(RuntimeError, match="Failed to create data mover workload"):
             executor.move_data(token="test_token", project_id="proj_id", cluster_id="cluster_id")
 
+    @patch("time.sleep")
     @patch.object(DGXCloudExecutor, "create_data_mover_workload")
     @patch.object(DGXCloudExecutor, "status")
-    def test_move_data_failed(self, mock_status, mock_create):
+    def test_move_data_failed(self, mock_status, mock_create, mock_sleep):
         mock_response = MagicMock()
         mock_response.status_code = 200
         mock_response.json.return_value = {"workloadId": "job123", "actualPhase": "Pending"}