class SlurmctldCharm(CharmBase):
slurm_instance_manager_cls = SlurmSnapInstanceManager
def __init__(self, *args):
super().__init__(*args)
self.dbd_requires = HostPortRequires(self, "slurmdbd-host-port")
self.fw_adapter = FrameworkAdapter(self.framework)
self.slurm_snap = self.slurm_instance_manager_cls(self, "slurmdctld")
self.munge = MungeProvides(self, "munge")
event_handler_bindings = {
self.on.install:
self._on_install,
self.dbd_requires.on.host_port_available:
self._on_dbd_host_port_available,
}
for event, handler in event_handler_bindings.items():
self.fw_adapter.observe(event, handler)
def _on_install(self, event):
handle_install(
event,
self.fw_adapter,
self.slurm_snap,
)
def _on_dbd_host_port_available(self, event):
handle_dbd_host_port_available(
event,
self.fw_adapter,
)
def __init__(self, charm, relation_name):
super().__init__(charm, relation_name)
self.fw_adapter = FrameworkAdapter(self.framework)
self.relation_name = relation_name
self.fw_adapter.observe(charm.on[relation_name].relation_changed,
self.on_relation_changed)
class SlurmdbdCharm(CharmBase):
_state = StoredState()
slurm_instance_manager_cls = SlurmSnapInstanceManager
def __init__(self, *args):
super().__init__(*args)
self._state.set_default(db_info=None)
self._state.set_default(munge_key=None)
# provides host port to slurmctld
self.dbd_provides = HostPortProvides(self, "slurmdbd-host-port", f'{socket.gethostname()}', '6819')
self.slurm_snap = self.slurm_instance_manager_cls(self, "slurmdbd")
self.fw_adapter = FrameworkAdapter(self.framework)
self.db = MySQLClient(self, "db")
self.munge = MungeRequires(self, "munge")
event_handler_bindings = {
self.db.on.database_available: self._on_database_available,
self.munge.on.munge_available: self._on_munge_available,
self.on.install: self._on_install,
self.on.start: self._on_start,
}
for event, handler in event_handler_bindings.items():
self.fw_adapter.observe(event, handler)
def _on_install(self, event):
handle_install(
event,
self.fw_adapter,
self.slurm_snap,
self.dbd_provides,
)
def _on_start(self, event):
handle_start(
event,
self.fw_adapter,
self.slurm_snap,
self._state,
)
def _on_database_available(self, event):
handle_database_available(
event,
self.fw_adapter,
self._state,
)
def _on_munge_available(self, event):
handle_munge_available(
event,
self.fw_adapter,
self.slurm_snap,
self._state,
)
def test__am_i_leader__returns_true_if_unit_is_leader(self):
# Setup
mock_framework = create_autospec(self.create_framework(),
spec_set=True)
mock_framework.model.unit.is_leader.return_value = True
# Exercise
adapter = FrameworkAdapter(mock_framework)
is_leader = adapter.am_i_leader()
# Assert
assert is_leader
class Charm(CharmBase):
_stored = StoredState()
def __init__(self, *args):
super().__init__(*args)
# Abstract out framework and friends so that this object is not
# too tightly coupled with the underlying framework's implementation.
# From this point forward, our Charm object will only interact with the
# adapter and not directly with the framework.
self.fw_adapter = FrameworkAdapter(self.framework)
self.prometheus = PrometheusInterface(self, 'http-api')
self.alertmanager = AlertManagerInterface(self, 'alertmanager')
# Bind event handlers to events
event_handler_bindings = {
self.on.start: self.on_start,
self.on.config_changed: self.on_config_changed,
self.on.upgrade_charm: self.on_upgrade,
self.on.stop: self.on_stop,
self.alertmanager.on.new_relation:
self.on_new_alertmanager_relation
}
for event, handler in event_handler_bindings.items():
self.fw_adapter.observe(event, handler)
self._stored.set_default(is_started=False)
# DELEGATORS
# These delegators exist to decouple the actual handlers from the
# underlying framework which has some very specific requirements that
# do not always apply to every event. For instance if we were to add
# an interface in our initializer, we would be forced to write unit
# tests that mock out that object even for handlers that do not need
# it. This hard coupling results in verbose tests that contain unused
# mocks. These tests tend to be hard to follow. To counter that, the
# logic is moved away from this class.
def on_config_changed(self, event):
on_config_changed_handler(event, self.fw_adapter, self._stored)
def on_new_alertmanager_relation(self, event):
on_new_alertmanager_relation_handler(event, self.fw_adapter)
def on_start(self, event):
on_start_handler(event, self.fw_adapter)
def on_upgrade(self, event):
on_upgrade_handler(event, self.fw_adapter)
def on_stop(self, event):
on_stop_handler(event, self.fw_adapter)
def test__get_app_name__returns_the_name_of_the_app(self):
# Setup
mock_framework = create_autospec(self.create_framework(),
spec_set=True)
mock_app_name = str(uuid4())
mock_framework.model.app.name = mock_app_name
# Exercise
adapter = FrameworkAdapter(mock_framework)
app_name = adapter.get_app_name()
# Assert
assert app_name == mock_app_name
def test__get_config__returns_a_value_given_a_key(self):
# Setup
mock_framework = create_autospec(self.create_framework(),
spec_set=True)
mock_key = str(uuid4())
mock_value = str(uuid4())
mock_framework.model.config = {mock_key: mock_value}
# Exercise
adapter = FrameworkAdapter(mock_framework)
value = adapter.get_config(mock_key)
# Assert
assert value == mock_value
class PrometheusInterface(Object):
on = PrometheusEvents()
def __init__(self, charm, relation_name):
super().__init__(charm, relation_name)
self.fw_adapter = FrameworkAdapter(self.framework)
self.relation_name = relation_name
self.fw_adapter.observe(charm.on[relation_name].relation_changed,
self.on_relation_changed)
def on_relation_changed(self, event):
logger.debug("Emitting new_prom_rel event")
self.on.new_prom_rel.emit()
logger.debug("Done emitting new_prom_rel_event")
def __init__(self, charm, key):
super().__init__(charm, key)
self.snap_mode = key
self.fw_adapter = FrameworkAdapter(self.framework)
# Set the template and config file paths based on the snap.mode.
# Throw an exception if initialized with an unsupported snap.mode.
if self.snap_mode == "slurmdbd":
self.slurm_config_yaml = self.SLURM_CONFIGURATOR_TEMPLATES_DIR / 'slurmdbd.yaml'
self.slurm_config_template = self.TEMPLATE_DIR / 'slurmdbd.yaml.tmpl'
elif self.snap_mode in ["slurmd", "slurmrestd", "slurmctld"]:
self.slurm_config_yaml = self.SLURM_CONFIGURATOR_TEMPLATES_DIR / 'slurm.yaml'
self.slurm_config_template = self.TEMPLATE_DIR / 'slurm.yaml.tmpl'
else:
self.logger.error(
f"Slurm component not supported: {self.snap_mode}",
exc_info=True)
class AlertManagerInterface(Object):
on = AlertManagerEvents()
def __init__(self, charm, relation_name):
super().__init__(charm, relation_name)
self.fw_adapter = FrameworkAdapter(self.framework)
self.relation_name = relation_name
self.fw_adapter.observe(charm.on[relation_name].relation_changed,
self.on_relation_changed)
def on_relation_changed(self, event):
remote_data = event.relation.data[event.unit]
logging.debug("Received remote_data: {}".format(dict(remote_data)))
logger.debug("Emitting new_relation event")
self.on.new_relation.emit(remote_data)
def test__get_config__returns_the_config_dict_when_no_key_given(self):
# Setup
mock_framework = create_autospec(self.create_framework(),
spec_set=True)
mock_config = {
str(uuid4()): str(uuid4()),
str(uuid4()): str(uuid4()),
str(uuid4()): str(uuid4()),
str(uuid4()): str(uuid4()),
}
mock_framework.model.config = mock_config
# Exercise
adapter = FrameworkAdapter(mock_framework)
config = adapter.get_config()
# Assert
assert config == mock_config
def test__get_image_meta__returns_an_image_meta_object(
self, mock_fetch_image_meta_func):
# Setup
framework = self.create_framework()
mock_framework = create_autospec(framework, spec_set=True)
mock_framework.model = create_autospec(framework.model, spec_set=True)
image_name = uuid4()
# Exercise
adapter = FrameworkAdapter(mock_framework)
image_meta = adapter.get_image_meta(image_name)
# Assert
assert mock_fetch_image_meta_func.call_count == 1
assert mock_fetch_image_meta_func.call_args == \
call(image_name, mock_framework.model.resources)
assert image_meta == mock_fetch_image_meta_func.return_value
class PrometheusInterface(Object):
on = PrometheusEvents()
def __init__(self, charm, relation_name):
super().__init__(charm, relation_name)
self.fw = FrameworkAdapter(self.framework)
self.relation_name = relation_name
self.fw.observe(charm.on[relation_name].relation_joined,
self.on_relation_joined)
def render_relation_data(self):
logging.debug('render-relation-data in')
for relation in self.model.relations[self.relation_name]:
relation.data[self.model.unit]['prometheus-port'] = \
str(PROMETHEUS_ADVERTISED_PORT)
logging.debug('render-relation-data out')
def on_relation_joined(self, event):
logging.debug("on-joined; emit new-client")
self.on.new_client.emit(event.relation)
self.render_relation_data()
def __init__(self, *args):
super().__init__(*args)
# Abstract out framework and friends so that this object is not
# too tightly coupled with the underlying framework's implementation.
# From this point forward, our Charm object will only interact with the
# adapter and not directly with the framework.
self.fw_adapter = FrameworkAdapter(self.framework)
self.prometheus = PrometheusInterface(self, 'http-api')
self.alertmanager = AlertManagerInterface(self, 'alertmanager')
# Bind event handlers to events
event_handler_bindings = {
self.on.start: self.on_start,
self.on.config_changed: self.on_config_changed,
self.on.upgrade_charm: self.on_upgrade,
self.on.stop: self.on_stop,
self.alertmanager.on.new_relation:
self.on_new_alertmanager_relation
}
for event, handler in event_handler_bindings.items():
self.fw_adapter.observe(event, handler)
self._stored.set_default(is_started=False)
def __init__(self, *args):
super().__init__(*args)
# Abstract out framework and friends so that this object is not
# too tightly coupled with the underlying framework's implementation.
# From this point forward, our Charm object will only interact with the
# adapter and not directly with the framework.
self.fw_adapter = FrameworkAdapter(self.framework)
self.prom_relation_name = 'prometheus'
self.prom_interface = \
PrometheusInterface(self, self.prom_relation_name)
# Bind event handlers to events
event_handler_bindings = {
self.on.start: self.on_start,
self.on.config_changed: self.on_config_changed,
self.on.upgrade_charm: self.on_upgrade,
self.on.stop: self.on_stop,
self.prom_interface.on.new_prom_rel: self.on_new_prom_rel
}
for event, handler in event_handler_bindings.items():
self.fw_adapter.observe(event, handler)
class SlurmSnapInstanceManager(Object):
"""
responsible for installing the slurm_snap, connecting to network, and
setting the snap mode
"""
_stored = StoredState()
on = SlurmSnapInstanceManagerEvents()
logger = logging.getLogger()
MUNGE_KEY_PATH = Path("/var/snap/slurm/common/etc/munge/munge.key")
SLURM_CONFIGURATOR_TEMPLATES_DIR = Path(
"/var/snap/slurm/common/etc/slurm-configurator")
TEMPLATE_DIR = Path(f"{os.getcwd()}/templates")
def __init__(self, charm, key):
super().__init__(charm, key)
self.snap_mode = key
self.fw_adapter = FrameworkAdapter(self.framework)
# Set the template and config file paths based on the snap.mode.
# Throw an exception if initialized with an unsupported snap.mode.
if self.snap_mode == "slurmdbd":
self.slurm_config_yaml = self.SLURM_CONFIGURATOR_TEMPLATES_DIR / 'slurmdbd.yaml'
self.slurm_config_template = self.TEMPLATE_DIR / 'slurmdbd.yaml.tmpl'
elif self.snap_mode in ["slurmd", "slurmrestd", "slurmctld"]:
self.slurm_config_yaml = self.SLURM_CONFIGURATOR_TEMPLATES_DIR / 'slurm.yaml'
self.slurm_config_template = self.TEMPLATE_DIR / 'slurm.yaml.tmpl'
else:
self.logger.error(
f"Slurm component not supported: {self.snap_mode}",
exc_info=True)
@property
def _hostname(self):
return socket.gethostname().split(".")[0]
def set_snap_mode(self):
"""Set the snap mode, thorw an exception if it fails.
"""
try:
subprocess.call([
"snap",
"set",
"slurm",
f"snap.mode={self.snap_mode}",
])
except subprocess.CalledProcessError as e:
self.logger.error(
f"Setting the snap.mode failed. snap.mode={self.snap_mode} - {e}",
exc_info=True)
def install(self):
self._install_snap()
self._snap_connect()
def _snap_connect(self, slot=None):
connect_commands = [
["snap", "connect", "slurm:network-control"],
["snap", "connect", "slurm:system-observe"],
["snap", "connect", "slurm:hardware-observe"],
]
for connect_command in connect_commands:
if slot:
connect_command.append(slot)
try:
subprocess.call(connect_command)
except subprocess.CalledProcessError as e:
self.logger.error(
f"Could not connect snap interface: {e}",
exc_info=True,
)
def _install_snap(self):
snap_install_cmd = ["snap", "install"]
resource_path = None
try:
resource_path = self.model.resources.fetch('slurm')
except ModelError as e:
logger.error(
f"Resource could not be found when executing: {e}",
exc_info=True,
)
if resource_path:
snap_install_cmd.append(resource_path)
snap_install_cmd.append("--dangerous")
else:
snap_store_channel = self.fw_adapter.get_config(
"snap-store-channel")
snap_install_cmd.append("slurm")
snap_install_cmd.append(f"--{snap_store_channel}")
try:
subprocess.call(snap_install_cmd)
except subprocess.CalledProcessError as e:
logger.error(
f"Could not install the slurm snap using the command: {e}",
exc_info=True)
def write_munge_key(self, munge_key):
munge = b64decode(munge_key.encode())
self.MUNGE_KEY_PATH.write_bytes(munge)
def write_config(self, context):
ctxt = {}
source = self.slurm_config_template
target = self.slurm_config_yaml
if not type(context) == dict:
self.framework.set_unit_status(
MaintenanceStatus("context not of type dict"))
return
else:
ctxt = {**{"hostname": self._hostname}, **context}
if not source.exists():
# raise Exception(f"Source config {source} does not exist - Please debug.")
self.framework.set_unit_status(
MaintenanceStatus("source doesn't exist"))
if target.exists():
target.unlink()
target.write_text(source.read_text().format(**ctxt))
