Skip to content

Commit

Permalink
Add selective shutdown
Browse files Browse the repository at this point in the history
New -s/--service flag for shutdown, to select a single service for shutdown if required
  • Loading branch information
vaibhav-dhawan committed Nov 22, 2023
1 parent 028d9e2 commit 7ecf85f
Show file tree
Hide file tree
Showing 2 changed files with 32 additions and 1 deletion.
28 changes: 27 additions & 1 deletion domino_maintenance_mode/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -102,6 +102,18 @@ async def _async_snapshot(output, **kwargs):
json.dump(state, output)


def validate_services(ctx, param, value):
if not value:
return None
elif value not in __get_execution_interfaces().keys():
raise click.BadParameter(
"Services must be one of "
f"{list(__get_execution_interfaces().keys())}"
)
else:
return value


@click.command()
@click.argument("snapshot", type=click.File("r"))
@click.option(
Expand Down Expand Up @@ -138,17 +150,31 @@ async def _async_snapshot(output, **kwargs):
default=600,
help="Amount of time to wait for executions to complete.",
)
@click.option(
"-s",
"--service",
type=str,
help="(Optional) Service to shutdown. Options are: "
f"{list(__get_execution_interfaces().keys())}",
callback=validate_services,
)
def shutdown(snapshot, **kwargs):
"""Stop running Apps, Model APIs, Durable Workspaces, and Scheduled Jobs.
SNAPSHOT : The path to snapshot output from 'dmm snapshot'.
"""
state = __load_state(snapshot)
manager = Manager(**kwargs)
for interface in __get_execution_interfaces().values():
if manager.get_service():
interface = __get_execution_interfaces()[manager.get_service()]
executions = state[interface.singular()]
if len(executions) > 0:
manager.stop(interface, executions)
else:
for interface in __get_execution_interfaces().values():
executions = state[interface.singular()]
if len(executions) > 0:
manager.stop(interface, executions)


cli.add_command(shutdown)
Expand Down
5 changes: 5 additions & 0 deletions domino_maintenance_mode/manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ class Manager:

def __init__(
self,
service: str = None,
batch_size: int = 5,
batch_interval_s: int = 5,
max_failures: int = 5,
Expand All @@ -37,6 +38,10 @@ def __init__(
self.batch_interval_s = batch_interval_s
self.grace_period_s = grace_period_s
self.max_failures = max_failures
self.service = service

def get_service(self):
return self.service

def stop(self, interface: ExecutionInterface, executions: List[Execution]):
self.__toggle_executions(
Expand Down

0 comments on commit 7ecf85f

Please sign in to comment.