(monarch/tools) Add force_restart argument to get_or_create command (meta-pytorch#783)

Summary:
Pull Request resolved: meta-pytorch#783

`force_restart=True` will kill the existing job and force it to create again. Useful when we want to redeploy the server to reflect any changes in the local conda/workspace that are not eligible for CodeSync.

NOTE: this is tangential to setting `workspace=None` which disables local conda AND workspace building into an ephemeral image.

Reviewed By: ahmadsharif1

Differential Revision: D79754696

fbshipit-source-id: 085aa2f1633d5b4ebac375a8494fc754a532d4a1

Author: kiukchung

python/monarch/tools/commands.py

@@ -265,6 +265,7 @@ async def get_or_create(
     name: str,
     config: Config,
     check_interval: timedelta = _5_SECONDS,
+    force_restart: bool = False,
 ) -> ServerSpec:
     """Waits for the server based on identity `name` in the scheduler specified in the `config`
     to be ready (e.g. RUNNING). If the server is not found then this function creates one
@@ -286,6 +287,7 @@ async def get_or_create(
         name: the name of the server (job) to get or create
         config: configs used to create the job if one does not exist
         check_interval: how often to poll the status of the job when waiting for it to be ready
+        force_restart: if True kills and re-creates the job even if one exists
 
     Returns: A `ServerSpec` containing information about either the existing or the newly
         created server.
@@ -322,6 +324,12 @@ async def get_or_create(
         return server_info
     else:
         print(f"{CYAN}Found existing job `{server_handle}` ready to serve.{ENDC}")
+
+        if force_restart:
+            print(f"{CYAN}force_restart=True, restarting `{server_handle}`.{ENDC}")
+            kill(server_handle)
+            server_info = await get_or_create(name, config, check_interval)
+
         return server_info
 
 

python/tests/tools/test_commands.py

@@ -24,6 +24,7 @@
 
 CMD_INFO = "monarch.tools.commands.info"
 CMD_CREATE = "monarch.tools.commands.create"
+CMD_KILL = "monarch.tools.commands.kill"
 
 
 class TestCommands(unittest.TestCase):
@@ -336,3 +337,45 @@ async def test_get_or_create_new_server_missing(
                 config=config,
                 check_interval=_5_MS,
             )
+
+    async def test_get_or_create_force_restart(self) -> None:
+        with mock.patch(
+            CMD_INFO,
+            side_effect=[
+                # -- state for slurm:///123
+                server(AppState.RUNNING, name="123"),
+                # -- force_restart kills the server
+                server(AppState.CANCELLED, name="123"),
+                # -- states for (new) slurm:///456
+                server(AppState.SUBMITTED, name="456"),
+                server(AppState.PENDING, name="456"),
+                server(AppState.RUNNING, name="456"),
+            ],
+        ) as mock_info, mock.patch(
+            CMD_CREATE, return_value="slurm:///456"
+        ) as mock_create, mock.patch(CMD_KILL) as mock_kill:
+            config = Config(
+                scheduler="slurm",
+                scheduler_args={},
+                appdef=defaults.component_fn("slurm")(),
+            )
+            server_info = await commands.get_or_create(
+                name="123",
+                config=config,
+                check_interval=_5_MS,
+                force_restart=True,
+            )
+
+            mock_create.called_once_with(config, "123")
+            mock_kill.assert_called_once_with("slurm:///123")
+            self.assertEqual(server_info.server_handle, "slurm:///456")
+            self.assertListEqual(
+                mock_info.call_args_list,
+                [
+                    mock.call("slurm:///123"),
+                    mock.call("slurm:///123"),
+                    mock.call("slurm:///456"),
+                    mock.call("slurm:///456"),
+                    mock.call("slurm:///456"),
+                ],
+            )