(monarch/tools) Add force_restart argument to get_or_create command

python/monarch/tools/colors.py

@@ -0,0 +1,25 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+
+import sys
+
+# only print colors if outputting directly to a terminal
+if not sys.stdout.closed and sys.stdout.isatty():
+    GREEN = "\033[32m"
+    BLUE = "\033[34m"
+    ORANGE = "\033[38:2:238:76:44m"
+    GRAY = "\033[2m"
+    CYAN = "\033[36m"
+    ENDC = "\033[0m"
+else:
+    GREEN = ""
+    ORANGE = ""
+    BLUE = ""
+    GRAY = ""
+    CYAN = ""
+    ENDC = ""

python/monarch/tools/commands.py

@@ -14,6 +14,8 @@
 from datetime import datetime, timedelta
 from typing import Any, Callable, Mapping, Optional, Union
 
+from monarch.tools.colors import CYAN, ENDC
+
 from monarch.tools.components.hyperactor import DEFAULT_NAME
 
 from monarch.tools.config import (  # @manual=//monarch/python/monarch/tools/config/meta:defaults
@@ -263,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
@@ -280,6 +283,12 @@ async def get_or_create(
         server_handle = get_or_create(name="my_job_name", config)
         server_info = info(server_handle)
 
+    Args:
+        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.
 
@@ -311,10 +320,16 @@ async def get_or_create(
                 f"the new server `{new_server_handle}` has {server_info.state}"
             )
 
-        print(f"\x1b[36mNew job `{new_server_handle}` is ready to serve. \x1b[0m")
+        print(f"{CYAN}New job `{new_server_handle}` is ready to serve.{ENDC}")
         return server_info
     else:
-        print(f"\x1b[36mFound existing job `{server_handle}` ready to serve. \x1b[0m")
+        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"),
+                ],
+            )