3 from __future__ import absolute_import, print_function
9 ComputeNodeSetupActor, ComputeNodeMonitorActor
10 from . import ComputeNodeShutdownActor as ShutdownActorBase
11 from . import ComputeNodeUpdateActor as UpdateActorBase
12 from .. import RetryMixin
14 class SlurmMixin(object):
15 SLURM_END_STATES = frozenset(['down\n', 'down*\n',
16 'drain\n', 'drain*\n',
18 SLURM_DRAIN_STATES = frozenset(['drain\n', 'drng\n'])
20 def _set_node_state(self, nodename, state, *args):
21 cmd = ['scontrol', 'update', 'NodeName=' + nodename,
24 subprocess.check_output(cmd)
26 def _get_slurm_state(self, nodename):
27 return subprocess.check_output(['sinfo', '--noheader', '-o', '%t', '-n', nodename])
30 class ComputeNodeShutdownActor(SlurmMixin, ShutdownActorBase):
32 arv_node = self._arvados_node()
35 return super(ComputeNodeShutdownActor, self).on_start()
38 self._nodename = arv_node['hostname']
39 self._logger.info("Draining SLURM node %s", self._nodename)
40 self._later.issue_slurm_drain()
42 @RetryMixin._retry((subprocess.CalledProcessError, OSError))
43 def cancel_shutdown(self, reason, try_resume=True):
45 if try_resume and self._get_slurm_state(self._nodename) in self.SLURM_DRAIN_STATES:
46 # Resume from "drng" or "drain"
47 self._set_node_state(self._nodename, 'RESUME')
49 # Node is in a state such as 'idle' or 'alloc' so don't
50 # try to resume it because that will just raise an error.
52 return super(ComputeNodeShutdownActor, self).cancel_shutdown(reason)
54 @RetryMixin._retry((subprocess.CalledProcessError, OSError))
55 def issue_slurm_drain(self):
56 if self.cancel_reason is not None:
59 self._set_node_state(self._nodename, 'DRAIN', 'Reason=Node Manager shutdown')
60 self._logger.info("Waiting for SLURM node %s to drain", self._nodename)
61 self._later.await_slurm_drain()
63 self._later.shutdown_node()
65 @RetryMixin._retry((subprocess.CalledProcessError, OSError))
66 def await_slurm_drain(self):
67 if self.cancel_reason is not None:
69 output = self._get_slurm_state(self._nodename)
70 if output in ("drng\n", "alloc\n", "drng*\n", "alloc*\n"):
71 self._timer.schedule(time.time() + 10,
72 self._later.await_slurm_drain)
73 elif output in ("idle\n"):
74 # Not in "drng" but idle, don't shut down
75 self.cancel_shutdown("slurm state is %s" % output.strip(), try_resume=False)
78 self._later.shutdown_node()
80 def _destroy_node(self):
82 self._set_node_state(self._nodename, 'DOWN', 'Reason=Node Manager shutdown')
83 super(ComputeNodeShutdownActor, self)._destroy_node()
86 class ComputeNodeUpdateActor(UpdateActorBase):
87 def sync_node(self, cloud_node, arvados_node):
88 if arvados_node.get("hostname"):
90 subprocess.check_output(['scontrol', 'update', 'NodeName=' + arvados_node["hostname"], 'Weight=%i' % int(cloud_node.size.price * 1000)])
92 self._logger.error("Unable to set slurm node weight.", exc_info=True)
93 return super(ComputeNodeUpdateActor, self).sync_node(cloud_node, arvados_node)