update DFS

csst_dag/dfs.py

@@ -3,26 +3,46 @@ import toml
 import os
 from csst_dfs_client import plan, level0
 from astropy.table import Table
+import socket
 
 CONFIG = toml.load(os.path.join(os.path.dirname(__file__), "config", "config.toml"))
 
 
+def check_port(ip, port, timeout=3):
+    """
+    # # 示例：检查 192.168.1.1 的 80 端口是否开放
+    # print(check_port("192.168.1.1", 80))  # True/False
+    """
+    try:
+        # 创建 Socket 连接（TCP）
+        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        sock.settimeout(timeout)  # 设置超时时间（秒）
+        # 尝试连接
+        result = sock.connect_ex((ip, port))
+        # 返回状态
+        if result == 0:  # 0 表示成功
+            return True
+        else:
+            return False
+    except Exception as e:
+        print(f"Error: {e}")
+        return False
+    finally:
+        sock.close()  # 确保关闭连接
+
+
 class DFS:
     def __init__(self, location=None):
         # try each location
         print("Test all locations:")
         status_list = []
         for loc in CONFIG.keys():
-            # print("Setting DFS config:")
-            for k, v in CONFIG[loc]["dfs"].items():
-                os.environ.setdefault(k, str(v))
-            # print("Setting redis config:")
-            redis = Redis(location=loc)
-            this_dfs_status = plan.find(instrument="").success
-            try:
-                this_redis_status = redis.ping()
-            except BaseException:
-                this_redis_status = False
+            dfs_ip = CONFIG[loc]["dfs"]["CSST_DFS_GATEWAY"].split(":")[0]
+            dfs_port = int(CONFIG[loc]["dfs"]["CSST_DFS_GATEWAY"].split(":")[1])
+            redis_ip = CONFIG[loc]["redis"]["host"]
+            redis_port = CONFIG[loc]["redis"]["port"]
+            this_dfs_status = check_port(dfs_ip, dfs_port, timeout=1)
+            this_redis_status = check_port(redis_ip, redis_port, timeout=1)
             this_status = dict(
                 location=loc,
                 status=this_dfs_status and this_redis_status,
@@ -30,6 +50,7 @@ class DFS:
                 redis=this_redis_status,
             )
             status_list.append(this_status)
+            print(this_status)
         status_table = Table(status_list)
         print(status_table)
 
@@ -49,7 +70,7 @@ class DFS:
 
 
 class Redis(redis.Redis):
-    def __init__(self, location="naoc"):
+    def __init__(self, location="naoc", **kwargs):
         if location not in CONFIG.keys():
             print("Available redis locations: ", list(CONFIG.keys()))
             raise ValueError(f"Unknown redis location: {location}")
@@ -58,6 +79,7 @@ class Redis(redis.Redis):
             port=CONFIG[location]["redis"]["port"],
             db=CONFIG[location]["redis"]["db"],
             password=CONFIG[location]["redis"]["password"],
+            **kwargs,
         )
         self.qname = password = CONFIG[location]["redis"]["qname"]
         self.config = CONFIG[location]["redis"]