From 4a6cd8e8a0ab437eec0b23ca329c495c2401c1bf Mon Sep 17 00:00:00 2001
From: Newnius <newnius.cn@gmail.com>
Date: Wed, 24 Apr 2019 15:14:01 +0800
Subject: [PATCH] update agent, add more metrics

---
 Dockerfile   |  2 +-
 README.md    | 28 +++++++++++++++++-----------
 bootstrap.sh |  3 +++
 executor.py  | 20 ++++++++++++++++++++
 monitor.py   | 15 ++++++++++++++-
 test.py      |  3 ++-
 6 files changed, 57 insertions(+), 14 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index be21dff..105a2d6 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -2,7 +2,7 @@ FROM quickdeploy/yao-python3
 
 MAINTAINER Newnius <newnius.cn@gmail.com>
 
-RUN pip3 install docker kafka
+RUN pip3 install docker kafka psutil
 
 ADD bootstrap.sh /etc/bootstrap.sh
 
diff --git a/README.md b/README.md
index e4e62f4..8f448c8 100644
--- a/README.md
+++ b/README.md
@@ -1,31 +1,37 @@
 # YAO-agent
 
-
-```bash
-bin/kafka-topics.sh \
-	--describe \
-	--zookeeper zookeeper_node1:2181,zookeeper_node2:2181,zookeeper_node3:2181 \
-	--topic yao
-```
-
 ```bash
 bin/kafka-topics.sh \
 	--create \
-	--zookeeper zookeeper_node1:2181,zookeeper_node2:2181,zookeeper_node3:2181 \
+	--zookeeper zookeeper-node1:2181,zookeeper-node2:2181,zookeeper-node3:2181 \
 	--replication-factor 3 \
 	--partitions 1 \
 	--topic yao
 ```
 
+```bash
+bin/kafka-topics.sh \
+	--describe \
+	--zookeeper zookeeper-node1:2181,zookeeper-node2:2181,zookeeper-node3:2181 \
+	--topic yao
+```
+
 ```bash
 bin/kafka-console-consumer.sh \
-	--bootstrap-server kafka_node1:9091,kafka_node2:9092,kafka_node3:9093 \
+	--bootstrap-server kafka-node1:9092,kafka-node2:9092,kafka-node3:9092 \
 	--topic yao \
 	--from-beginning
 ```
 
 ```bash
 bin/kafka-console-producer.sh \
-	--broker-list kafka_node1:9091,kafka_node2:9092,kafka_node3:9093 \
+	--broker-list kafka-node1:9092,kafka-node2:9092,kafka-node3:9092 \
+	--topic yao
+```
+
+```bash
+bin/kafka-topics.sh \
+	--delete \
+	--zookeeper zookeeper-node1:2181,zookeeper-node2:2181,zookeeper-node3:2181 \
 	--topic yao
 ```
\ No newline at end of file
diff --git a/bootstrap.sh b/bootstrap.sh
index 00077ad..d57b604 100755
--- a/bootstrap.sh
+++ b/bootstrap.sh
@@ -2,6 +2,9 @@
 
 # TODO: monitor the processes
 
+# run nvidia-smi in background to speed up the query and reduce CPU load (why?)
+nvidia-smi daemon
+
 python3 /root/monitor.py &
 
 python3 /root/executor.py &
diff --git a/executor.py b/executor.py
index 3ee7873..b619583 100644
--- a/executor.py
+++ b/executor.py
@@ -114,6 +114,26 @@ class MyHandler(BaseHTTPRequestHandler):
 			container.stop()
 			msg = {"code": 0}
 
+			self.send_response(200)
+			self.send_header('Content-type', 'application/json')
+			self.end_headers()
+			self.wfile.write(bytes(json.dumps(msg), "utf-8"))
+
+		elif self.path == "/remove":
+			form = cgi.FieldStorage(
+				fp=self.rfile,
+				headers=self.headers,
+				environ={
+					'REQUEST_METHOD': 'POST',
+					'CONTENT_TYPE': self.headers['Content-Type'],
+				})
+			container_id = form["id"].value
+
+			client = docker.from_env()
+			container = client.containers.get(container_id)
+			container.remove(force=True)
+			msg = {"code": 0}
+
 			self.send_response(200)
 			self.send_header('Content-type', 'application/json')
 			self.end_headers()
diff --git a/monitor.py b/monitor.py
index 0eaca8d..069e0df 100644
--- a/monitor.py
+++ b/monitor.py
@@ -5,6 +5,9 @@ import json
 from xml.dom.minidom import parse
 import xml.dom.minidom
 from kafka import KafkaProducer
+import multiprocessing
+import psutil
+import math
 
 ClientID = os.getenv('ClientID', 1)
 ClientHost = os.getenv('ClientHost', "localhost")
@@ -71,7 +74,17 @@ def report_msg():
 
 		stats.append(stat)
 
-	post_fields = {'id': ClientID, 'host': ClientHost, 'status': stats}
+	mem = psutil.virtual_memory()
+
+	post_fields = {
+		'id': ClientID,
+		'host': ClientHost,
+		'status': stats,
+		'cpu_num': multiprocessing.cpu_count(),
+		'cpu_load': os.getloadavg()[0],
+		'mem_total': math.floor(mem.total / (1024. ** 3)),
+		'mem_available': math.floor(mem.available / (1024. ** 3))
+	}
 	data = json.dumps(post_fields)
 
 	producer = KafkaProducer(bootstrap_servers=KafkaBrokers)
diff --git a/test.py b/test.py
index f01073d..0327dee 100644
--- a/test.py
+++ b/test.py
@@ -36,7 +36,7 @@ def get_status(id):
 	container = client.containers.list(all=True, filters={'id': id})
 	status = {}
 	if len(container) > 0:
-		container= container[0]
+		container = container[0]
 		status['id'] = container.short_id
 		status['image'] = container.attrs['Config']['Image']
 		status['image_digest'] = container.attrs['Image']
@@ -47,6 +47,7 @@ def get_status(id):
 		if status['command'] is not None:
 			status['command'] = ' '.join(container.attrs['Config']['Cmd'])
 	print(status)
+	print(container.attrs)
 
 
 get_status('')