diff --git a/src/slurm_autoscale_tfe/__init__.py b/src/slurm_autoscale_tfe/__init__.py index 6cd5890..4e75e0a 100644 --- a/src/slurm_autoscale_tfe/__init__.py +++ b/src/slurm_autoscale_tfe/__init__.py @@ -26,12 +26,15 @@ NODE_STATE_REGEX = re.compile(r"^NodeName=([a-z0-9-]*).*State=([A-Z_+]*).*$") DOWN_FLAG_SET = frozenset(["DOWN", "POWER_DOWN", "POWERED_DOWN", "POWERING_DOWN"]) -INSTANCE_TYPES = frozenset([ - "aws_instance", - "azurerm_linux_virtual_machine", - "google_compute_instance", - "openstack_compute_instance_v2", -]) +INSTANCE_TYPES = frozenset( + [ + "aws_instance", + "azurerm_linux_virtual_machine", + "google_compute_instance", + "openstack_compute_instance_v2", + ] +) + class AutoscaleException(Exception): """Raised when something bad happened in autoscale main""" @@ -66,7 +69,6 @@ def change_host_state(hostlist, state, reason=None): ) - def resume(hostlist=sys.argv[-1]): """Issue a request to Terraform cloud to power up the instances listed in hostlist. @@ -92,6 +94,7 @@ def suspend(hostlist=sys.argv[-1]): return 1 return 0 + def resume_fail(hostlist=sys.argv[-1]): """Issue a request to Terraform cloud to power down the instances listed in hostlist. @@ -106,8 +109,7 @@ def resume_fail(hostlist=sys.argv[-1]): def connect_tfe_client(): - """Return a TFE client object using environment variables for authentication - """ + """Return a TFE client object using environment variables for authentication""" if "TFE_TOKEN" not in environ: raise AutoscaleException( f"{sys.argv[0]} requires environment variable TFE_TOKEN" @@ -131,8 +133,7 @@ def connect_tfe_client(): def get_pool_from_tfe(tfe_client): - """Retrieve id and content of POOL variable from Terraform cloud - """ + """Retrieve id and content of POOL variable from Terraform cloud""" try: tfe_var = tfe_client.fetch_variable(POOL_VAR) except Timeout as exc: @@ -149,6 +150,7 @@ def get_pool_from_tfe(tfe_client): return tfe_var["id"], frozenset(tfe_var["value"]) return tfe_var["id"], frozenset() + def get_instances_from_tfe(tfe_client): """Return all names of instances that are created in Terraform Cloud state.""" try: @@ -163,6 +165,7 @@ def get_instances_from_tfe(tfe_client): address_prefix = resource["attributes"]["address"].split("[")[0] return frozenset(instances), address_prefix + def main(command, set_op, hostlist): """Issue a request to Terraform cloud to modify the pool variable of the workspace indicated by TFE_WORKSPACE environment variable using the operation @@ -192,7 +195,7 @@ def main(command, set_op, hostlist): try: run_id = tfe_client.apply( f"Slurm {command.value} {hostlist}".strip(), - targets=[f'module.{address_prefix}["{hostname}"]' for hostname in hosts] + targets=[f'module.{address_prefix}["{hostname}"]' for hostname in hosts], ) except Timeout as exc: raise AutoscaleException("Connection to Terraform cloud timeout (5s)") from exc