Skip to content
Snippets Groups Projects
Commit 7ec758f5 authored by Bo-Chun Chen's avatar Bo-Chun Chen
Browse files

Rewrite slurm account agent

parent a0911932
No related branches found
No related tags found
2 merge requests!23Feat resolve uid gid,!18Redesign RabbitMQ
......@@ -6,60 +6,40 @@ import socket
import subprocess
import time
import json
from rc_rmq import RCRMQ
hostname = socket.gethostname().split(".", 1)[0]
connect_host = rcfg.Server if hostname != rcfg.Server else "localhost"
queue_name = "slurm_add_account"
duration = 2
task = 'slurm_account'
# Set up credentials to connect to RabbitMQ server
credentials = pika.PlainCredentials(rcfg.User, rcfg.Password)
parameters = pika.ConnectionParameters(connect_host,
rcfg.Port,
rcfg.VHost,
credentials)
# Establish connection to RabbitMQ server
connection = pika.BlockingConnection(parameters)
channel = connection.channel()
print("connection established. Listening for messages:")
# create exchange to pass messages
channel.exchange_declare(exchange=rcfg.Exchange, exchange_type='direct')
# creates a random name for the newly generated queue
result = channel.queue_declare(queue=queue_name, exclusive=False)
channel.queue_bind(exchange=rcfg.Exchange, queue=queue_name, routing_key=queue_name)
# Instantiate rabbitmq object
rc_rmq = RCRMQ({'exchange': 'RegUsr', 'exchange_type': 'topic'})
def slurm_account_create(ch, method, properties, body):
msg = json.loads(body)
print("Message received {}".format(msg))
username = msg['username']
success = False
try:
subprocess.call(["sudo", "sacctmgr", "add", "account", username, "-i", "Descripition: Add user"])
subprocess.call(["sudo", "sacctmgr", "add", "user", username, "account="+username, "-i"])
print("SLURM account for user {} has been added".format(username))
success = True
except:
print("Failed to create user")
channel.basic_ack(delivery_tag=method.delivery_tag)
channel.basic_publish(exchange=rcfg.Exchange, routing_key=username, body=json.dumps(msg))
# ingest messages
channel.basic_consume(queue=queue_name, on_message_callback=slurm_account_create)
# initiate message ingestion
try:
channel.start_consuming()
except KeyboardInterrupt:
print("Disconnecting from broker.")
channel.stop_consuming()
connection.close()
ch.basic_ack(delivery_tag=method.delivery_tag)
# send confirm message
rc_rmq.publish_msg({
'routing_key': 'confirm.' + username,
'msg': {
'task': task,
'success': success
}
})
print("Start listening to queue: {}".format(task))
rc_rmq.start_consume({
'queue': task,
'routing_key': "create.*",
'cb': slurm_account_create
})
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment