Only ask for confirmation once (#39)

* Only ask for confirmation once Signed-off-by: Marc Romeijn <[email protected]> * Fix linting Signed-off-by: Marc Romeyn <[email protected]> * Fix failing test Signed-off-by: Marc Romeyn <[email protected]> * Detect torchrun Signed-off-by: Marc Romeyn <[email protected]> * Trying to fix failing test Signed-off-by: Marc Romeyn <[email protected]> * Fix failing tests Signed-off-by: Marc Romeyn <[email protected]> --------- Signed-off-by: Marc Romeijn <[email protected]> Signed-off-by: Marc Romeyn <[email protected]>
NVIDIA · Oct 2, 2024 · 0efe5d3 · 0efe5d3
1 parent 7217163
commit 0efe5d3
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 1 deletion.
diff --git a/src/nemo_run/cli/api.py b/src/nemo_run/cli/api.py
@@ -73,9 +73,12 @@
 DEFAULT_NAME = "default"
 EXECUTOR_CLASSES = [Executor, LocalExecutor, SkypilotExecutor, SlurmExecutor]
 PLUGIN_CLASSES = [Plugin, List[Plugin]]
+NEMORUN_SKIP_CONFIRMATION: Optional[bool] = None
 
 INCLUDE_WORKSPACE_FILE = os.environ.get("INCLUDE_WORKSPACE_FILE", "true").lower() == "true"
 
+logger = logging.getLogger(__name__)
+
 
 def entrypoint(
     fn: Optional[F] = None,
@@ -918,7 +921,24 @@ def _should_continue(self, skip_confirmation: bool) -> bool:
         Returns:
             bool: True if execution should continue, False otherwise.
         """
-        return skip_confirmation or typer.confirm("Continue?")
+        global NEMORUN_SKIP_CONFIRMATION
+
+        # If we're running under torchrun, always continue
+        if _is_torchrun():
+            logger.info("Detected torchrun environment. Skipping confirmation.")
+            return True
+
+        if NEMORUN_SKIP_CONFIRMATION is not None:
+            return NEMORUN_SKIP_CONFIRMATION
+
+        # If skip_confirmation is True or user confirms, continue
+        if skip_confirmation or typer.confirm("Continue?"):
+            NEMORUN_SKIP_CONFIRMATION = True
+            return True
+
+        # Otherwise, don't continue
+        NEMORUN_SKIP_CONFIRMATION = False
+        return False
 
     def parse_fn(self, fn: T, args: List[str], **default_kwargs) -> Partial[T]:
         """
@@ -1357,6 +1377,11 @@ class MissingRequiredOptionError(RunContextError):
     """Raised when a required option is missing."""
 
 
+def _is_torchrun() -> bool:
+    """Check if running under torchrun."""
+    return "WORLD_SIZE" in os.environ and int(os.environ["WORLD_SIZE"]) > 1
+
+
 if __name__ == "__main__":
     app = create_cli()
     app()
diff --git a/test/cli/test_api.py b/test/cli/test_api.py
@@ -26,6 +26,7 @@
 
 import nemo_run as run
 from nemo_run import cli, config
+from nemo_run.cli import api as cli_api
 from nemo_run.cli.api import Entrypoint, RunContext, create_cli
 
 _RUN_FACTORIES_ENTRYPOINT: str = """
@@ -187,6 +188,7 @@ def test_run_context_execute_task_with_confirmation_denied(
         self, mock_confirm, mock_run, mock_dryrun_fn, sample_function
     ):
         ctx = RunContext(name="test_run")
+        cli_api.NEMORUN_SKIP_CONFIRMATION = None
         ctx.cli_execute(sample_function, ["a=10", "b=hello"])
         mock_dryrun_fn.assert_called_once()
         mock_confirm.assert_called_once()

diff --git a/test/conftest.py b/test/conftest.py
@@ -16,6 +16,7 @@
 import os
 from typing import Any, Optional
 
+import pytest
 from invoke.config import Config
 from invoke.context import Context
 
@@ -33,3 +34,14 @@ def run(self, command: str, **kwargs: Any):
         kwargs["in_stream"] = False
         runner = self.config.runners.local(self)
         return self._run(runner, command, **kwargs)
+
+
+@pytest.fixture(autouse=True)
+def reset_nemorun_skip_confirmation():
+    from nemo_run.cli import api
+
+    """Reset NEMORUN_SKIP_CONFIRMATION to None before each test."""
+    api.NEMORUN_SKIP_CONFIRMATION = None
+    yield
+    # Optionally, reset after the test as well
+    api.NEMORUN_SKIP_CONFIRMATION = None