init refactor, add Scheduler

AlexCXC · AlexCXC · commit cf3123b6fca6 · 2025-08-28T18:32:23.000+08:00
diff --git a/.gitignore b/.gitignore
@@ -13,3 +13,5 @@ local_settings.py
 
 seatable-python-runner/
 seatable-python-runner.zip
+
+.python-version
diff --git a/scheduler/app/database/__init__.py b/scheduler/app/database/__init__.py
@@ -3,7 +3,7 @@
 
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy import create_engine
-from sqlalchemy.orm import sessionmaker
+from sqlalchemy.orm import sessionmaker, scoped_session
 
 DB_ROOT_USER = os.getenv("DB_ROOT_USER", "root")
 DB_ROOT_PASSWD = os.getenv("DB_ROOT_PASSWD")
@@ -37,4 +37,4 @@
 
 engine = create_engine(db_url, **db_kwargs)
 Base = declarative_base()
-DBSession = sessionmaker(bind=engine)
+DBSession = scoped_session(sessionmaker(bind=engine))
diff --git a/scheduler/app/faas_scheduler/models.py b/scheduler/app/faas_scheduler/models.py
@@ -31,6 +31,11 @@ class ScriptLog(Base):
     return_code = Column(Integer, nullable=True)
     output = Column(Text, nullable=True)
     operate_from = Column(String(255))
+    state = Column(String(10))
+
+    PENDING = 'pending'
+    RUNNING = 'running'
+    FINISHED = 'finished'
 
     def __init__(
         self,
@@ -40,6 +45,7 @@ def __init__(
         script_name,
         context_data,
         started_at,
+        state,
         operate_from=None,
     ):
         self.dtable_uuid = dtable_uuid
@@ -49,6 +55,7 @@ def __init__(
         self.context_data = context_data
         self.started_at = started_at
         self.operate_from = operate_from
+        self.state = state
 
     def to_dict(self):
         from faas_scheduler.utils import datetime_to_isoformat_timestr
@@ -68,6 +75,7 @@ def to_dict(self):
             "return_code": self.return_code,
             "output": self.output,
             "operate_from": self.operate_from,
+            "state": self.state
         }
 
 
diff --git a/scheduler/app/flask_server.py b/scheduler/app/flask_server.py
@@ -8,14 +8,11 @@
 from datetime import datetime
 from flask import Flask, request, make_response
 from gevent.pywsgi import WSGIServer
-from concurrent.futures import ThreadPoolExecutor
 
 from database import DBSession
 from faas_scheduler.utils import (
     check_auth_token,
-    run_script,
     get_script,
-    add_script,
     get_run_script_statistics_by_month,
     hook_update_script,
     can_run_task,
@@ -26,6 +23,7 @@
     uuid_str_to_32_chars,
     basic_log,
 )
+from .scheduler import scheduler
 
 
 basic_log("scheduler.log")
@@ -38,7 +36,11 @@
 app = Flask(__name__)
 
 logger = logging.getLogger(__name__)
-executor = ThreadPoolExecutor(max_workers=SCRIPT_WORKERS)
+
+
+@app.teardown_appcontext
+def shutdown_session(exception=None):
+    DBSession.remove()
 
 
 @app.route("/ping/", methods=["GET"])
@@ -72,8 +74,6 @@ def scripts_api():
     context_data = data.get("context_data")
     owner = data.get("owner")
     org_id = data.get("org_id")
-    script_url = data.get("script_url")
-    temp_api_token = data.get("temp_api_token")
     scripts_running_limit = data.get("scripts_running_limit", -1)
     operate_from = data.get("operate_from", "manualy")
     if not dtable_uuid or not script_name or not owner:
@@ -87,27 +87,16 @@ def scripts_api():
             owner, org_id, db_session, scripts_running_limit=scripts_running_limit
         ):
             return make_response(("The number of runs exceeds the limit"), 400)
-        script = add_script(
-            db_session,
-            dtable_uuid,
-            owner,
+        script_log = scheduler.add_script_log(
+            uuid_str_to_32_chars(dtable_uuid),
             org_id,
+            owner,
             script_name,
             context_data,
-            operate_from,
-        )
-        logger.debug("lets call the starter to fire up the runner...")
-        executor.submit(
-            run_script,
-            script.id,
-            dtable_uuid,
-            script_name,
-            script_url,
-            temp_api_token,
-            context_data,
+            operate_from
         )
 
-        return make_response(({"script_id": script.id}, 200))
+        return make_response(({"script_id": script_log.id}, 200))
     except Exception as e:
         logger.exception(e)
         return make_response(("Internal server error", 500))
@@ -143,15 +132,15 @@ def script_api(script_id):
         if dtable_uuid != script.dtable_uuid or script_name != script.script_name:
             return make_response(("Bad request", 400))
 
-        if SUB_PROCESS_TIMEOUT and isinstance(SUB_PROCESS_TIMEOUT, int):
-            now = datetime.now()
-            duration_seconds = (now - script.started_at).seconds
-            if duration_seconds > SUB_PROCESS_TIMEOUT:
-                script.success = False
-                script.return_code = -1
-                script.finished_at = now
-                script.output = TIMEOUT_OUTPUT
-                db_session.commit()
+        # if SUB_PROCESS_TIMEOUT and isinstance(SUB_PROCESS_TIMEOUT, int):
+        #     now = datetime.now()
+        #     duration_seconds = (now - script.started_at).seconds
+        #     if duration_seconds > SUB_PROCESS_TIMEOUT:
+        #         script.success = False
+        #         script.return_code = -1
+        #         script.finished_at = now
+        #         script.output = TIMEOUT_OUTPUT
+        #         db_session.commit()
 
         return make_response(({"script": script.to_dict()}, 200))
 
diff --git a/scheduler/app/scheduler.py b/scheduler/app/scheduler.py
@@ -1,56 +1,183 @@
+import json
+import logging
 import os
-import gc
 import time
-import logging
-from threading import Thread
+from threading import Lock, Thread
 
 from database import DBSession
+from faas_scheduler.models import ScriptLog
 from faas_scheduler.utils import (
-    check_and_set_tasks_timeout,
-    delete_log_after_days,
-    delete_statistics_after_days,
-    basic_log,
+    add_script,
+    run_script,
+    get_script_file,
+    hook_update_script
 )
 
-basic_log("scheduler.log")
+logger = logging.getLogger(__name__)
+
 
-SUB_PROCESS_TIMEOUT = int(
-    os.environ.get("PYTHON_PROCESS_TIMEOUT", 60 * 15)
-)  # 15 minutes
+class ScriptQueue:
 
-logger = logging.getLogger(__name__)
+    def __init__(self):
+        self.q = []  # a list of ScriptLog
+        self.script_logs_dict = {}  # a dict of {id: ScriptLog}
+        self.lock = Lock()
+        self.running_count = {}
+        # a dict of 
+        # {
+        #     "<team>": 0,
+        #     "<team>_<dtable_uuid>": 0,
+        #     "<team>_<dtable_uuid>_<script_name>": 0
+        # }
+        try:
+            run_limit_per_team = os.environ.get('RUN_LIMIT_PER_TEAM', 0)
+        except:
+            run_limit_per_team = 0
+        try:
+            run_limit_per_base = os.environ.get('RUN_LIMIT_PER_BASE', 0)
+        except:
+            run_limit_per_base = 0
+        try:
+            run_limit_per_script = os.environ.get('RUN_LIMIT_PER_SCRIPT', 0)
+        except:
+            run_limit_per_script = 0
+        self.config = {
+            'run_limit_per_team': run_limit_per_team,
+            'run_limit_per_base': run_limit_per_base,
+            'run_limit_per_script': run_limit_per_script
+        }
+
+    def can_run_script(self, script_log: ScriptLog):
+        if script_log.org_id != -1:
+            running_team_key = f'{script_log.org_id}'
+        else:
+            running_team_key = f'{script_log.owner}'
+        running_base_key = f'{running_team_key}_{script_log.dtable_uuid}'
+        running_script_key = f'{running_base_key}_{script_log.script_name}'
+
+        if self.config['run_limit_per_team'] > 0 and self.config['run_limit_per_team'] <= self.running_count.get(running_team_key, 0):
+            return False
+        if self.config['run_limit_per_base'] > 0 and self.config['run_limit_per_base'] <= self.running_count.get(running_base_key, 0):
+            return False
+        if self.config['run_limit_per_script'] > 0 and self.config['run_limit_per_script'] <= self.running_count.get(running_script_key, 0):
+            return False
+
+        return True
+
+    def add_script_log(self, script_log: ScriptLog):
+        with self.lock:
+            self.q.append(script_log)
+            self.script_logs_dict[script_log.id] = script_log
+
+    def get(self):
+        """get the first valid task from self.q
 
+        Return: an instance of ScriptTask or None
+        """
+        with self.lock:
+            return_task = None
 
-class FAASTaskTimeoutSetter(Thread):
+            index = 0
+            while index < len(self.q):
+                script_log = self.q[index]
+                if self.can_run_script(script_log):
+                    return_task = script_log
+                    self.q.pop(index)
+                    break
+                index += 1
+
+            return return_task
+
+    def script_done_callback(self, script_log: ScriptLog):
+        with self.lock:
+            if script_log.org_id != -1:
+                running_team_key = f'{script_log.org_id}'
+            else:
+                running_team_key = f'{script_log.owner}'
+            running_base_key = f'{running_team_key}_{script_log.dtable_uuid}'
+            running_script_key = f'{running_base_key}_{script_log.script_name}'
+            if running_team_key in self.running_count:
+                self.running_count[running_team_key] -= 1
+            if running_base_key in self.running_count:
+                self.running_count[running_base_key] -= 1
+            if running_script_key in self.running_count:
+                self.running_count[running_script_key] -= 1
+
+
+class Scheduelr:
 
     def __init__(self):
-        super(FAASTaskTimeoutSetter, self).__init__()
-        self.interval = 60 * 30  # every half an hour
-
-    def run(self):
-        if SUB_PROCESS_TIMEOUT and isinstance(SUB_PROCESS_TIMEOUT, int):
-            while True:
-                logger.info("Start automatic cleanup ...")
-                db_session = DBSession()
-                try:
-                    check_and_set_tasks_timeout(db_session)
-                except Exception as e:
-                    logger.exception("task cleaner error: %s", e)
-                finally:
-                    db_session.close()
-
-                # python garbage collection
-                logger.info("gc.collect: %s", str(gc.collect()))
-
-                # remove old script_logs and statistics
-                delete_log_after_days(db_session)
-                delete_statistics_after_days(db_session)
-
-                # sleep
-                logger.info("Sleep for %d seconds ...", self.interval)
-                time.sleep(self.interval)
-
-
-if __name__ == "__main__":
-    task_timeout_setter = FAASTaskTimeoutSetter()
-    task_timeout_setter.start()
+        self.script_queue = ScriptQueue()
+
+    def add_script_log(
+            self,
+            dtable_uuid,
+            org_id,
+            owner,
+            script_name,
+            context_data,
+            operate_from
+        ):
+        script_log = add_script(
+            DBSession(),
+            dtable_uuid,
+            owner,
+            org_id,
+            script_name,
+            context_data,
+            operate_from
+        )
+        self.script_queue.add_script_log(script_log)
+        return script_log
+
+    def schedule(self):
+        while True:
+            script_log = self.script_queue.get()
+            if not script_log:
+                time.sleep(0.5)
+            try:
+                script_file_info = get_script_file(script_log.dtable_uuid, script_log.script_name)
+                run_script(
+                    script_log.script_id,
+                    script_log.dtable_uuid,
+                    script_log.script_name,
+                    script_file_info['script_url'],
+                    script_file_info['temp_api_token'],
+                    json.loads(script_log.context_data)
+                )
+            except Exception as e:
+                logger.exception(f'run script: {script_log} error {e}')
+
+    def script_done_callback(
+            self,
+            script_id,
+            success,
+            return_code,
+            output,
+            spend_time
+        ):
+        script_log = self.script_queue.script_logs_dict.pop(script_id)
+        hook_update_script(
+            DBSession(),
+            script_id,
+            success,
+            return_code,
+            output,
+            spend_time
+        )
+        if not script_log:  # not counted in memory, only update db record
+            return
+
+    def load_pending_script_logs(self):
+        """load pending script logs, should be called only when server start
+        """
+        script_logs = DBSession.query(ScriptLog).filter_by(state=ScriptLog.PENDING).order_by(ScriptLog.id)
+        for script_log in script_logs:
+            self.script_queue.add_script_log(script_log)
+
+    def start(self):
+        self.load_pending_script_logs()
+        Thread(target=self.schedule, daemon=True).start()
+
+
+scheduler = Scheduelr()
diff --git a/scheduler/app/timeout_setter.py b/scheduler/app/timeout_setter.py
