2 # Copyright (C) The Arvados Authors. All rights reserved.
4 # SPDX-License-Identifier: AGPL-3.0
6 from __future__ import absolute_import, print_function
8 import subprocess32 as subprocess
11 from . import ComputeNodeMonitorActor
12 from . import ComputeNodeSetupActor as SetupActorBase
13 from . import ComputeNodeShutdownActor as ShutdownActorBase
14 from . import ComputeNodeUpdateActor as UpdateActorBase
15 from .. import RetryMixin
17 class SlurmMixin(object):
18 SLURM_END_STATES = frozenset(['down\n', 'down*\n',
19 'drain\n', 'drain*\n',
21 SLURM_DRAIN_STATES = frozenset(['drain\n', 'drng\n'])
23 def _update_slurm_node(self, nodename, updates):
24 cmd = ['scontrol', 'update', 'NodeName=' + nodename] + updates
26 subprocess.check_output(cmd)
29 "SLURM update %r failed", cmd, exc_info=True)
31 def _update_slurm_size_attrs(self, nodename, size):
32 self._update_slurm_node(nodename, [
33 'Weight=%i' % int(size.price * 1000),
34 'Features=instancetype=' + size.id,
37 def _get_slurm_state(self, nodename):
38 return subprocess.check_output(['sinfo', '--noheader', '-o', '%t', '-n', nodename])
41 class ComputeNodeSetupActor(SlurmMixin, SetupActorBase):
42 def create_cloud_node(self):
43 hostname = self.arvados_node.get("hostname")
45 self._update_slurm_size_attrs(hostname, self.cloud_size)
46 return super(ComputeNodeSetupActor, self).create_cloud_node()
49 class ComputeNodeShutdownActor(SlurmMixin, ShutdownActorBase):
51 arv_node = self._arvados_node()
54 return super(ComputeNodeShutdownActor, self).on_start()
57 self._nodename = arv_node['hostname']
58 self._logger.info("Draining SLURM node %s", self._nodename)
59 self._later.issue_slurm_drain()
61 @RetryMixin._retry((subprocess.CalledProcessError, OSError))
62 def cancel_shutdown(self, reason, try_resume=True):
64 if try_resume and self._get_slurm_state(self._nodename) in self.SLURM_DRAIN_STATES:
65 # Resume from "drng" or "drain"
66 self._update_slurm_node(self._nodename, ['State=RESUME'])
68 # Node is in a state such as 'idle' or 'alloc' so don't
69 # try to resume it because that will just raise an error.
71 return super(ComputeNodeShutdownActor, self).cancel_shutdown(reason)
73 @RetryMixin._retry((subprocess.CalledProcessError, OSError))
74 def issue_slurm_drain(self):
75 if self.cancel_reason is not None:
78 self._update_slurm_node(self._nodename, [
79 'State=DRAIN', 'Reason=Node Manager shutdown'])
80 self._logger.info("Waiting for SLURM node %s to drain", self._nodename)
81 self._later.await_slurm_drain()
83 self._later.shutdown_node()
85 @RetryMixin._retry((subprocess.CalledProcessError, OSError))
86 def await_slurm_drain(self):
87 if self.cancel_reason is not None:
89 output = self._get_slurm_state(self._nodename)
90 if output in ("drng\n", "alloc\n", "drng*\n", "alloc*\n"):
91 self._timer.schedule(time.time() + 10,
92 self._later.await_slurm_drain)
93 elif output in ("idle\n",):
94 # Not in "drng" but idle, don't shut down
95 self.cancel_shutdown("slurm state is %s" % output.strip(), try_resume=False)
98 self._later.shutdown_node()
100 def _destroy_node(self):
102 self._update_slurm_node(self._nodename, [
103 'State=DOWN', 'Reason=Node Manager shutdown'])
104 super(ComputeNodeShutdownActor, self)._destroy_node()
107 class ComputeNodeUpdateActor(SlurmMixin, UpdateActorBase):
108 def sync_node(self, cloud_node, arvados_node):
109 """Keep SLURM's node properties up to date."""
110 hostname = arvados_node.get("hostname")
111 features = arvados_node.get("slurm_node_features", "").split(",")
112 sizefeature = "instancetype=" + cloud_node.size.id
113 if hostname and sizefeature not in features:
114 # This probably means SLURM has restarted and lost our
115 # dynamically configured node weights and features.
116 self._update_slurm_size_attrs(hostname, cloud_node.size)
117 return super(ComputeNodeUpdateActor, self).sync_node(
118 cloud_node, arvados_node)