Support win32 for web-ui (#742)

swift/ui/llm_infer/runtime.py

@@ -1,5 +1,6 @@
 import collections
 import os.path
+import sys
 import time
 from datetime import datetime
 from typing import Dict, List, Tuple, Type
@@ -125,7 +126,7 @@ def update_log(cls):
     def wait(cls, task):
         if not task:
             return [None]
-        args = cls.parse_info_from_cmdline(task)
+        _, args = cls.parse_info_from_cmdline(task)
         log_file = args['log_file']
         offset = 0
         latest_data = ''
@@ -177,7 +178,7 @@ def get_all_ports():
                     ports.add(
                         int(
                             Runtime.parse_info_from_cmdline(
-                                Runtime.construct_running_task(proc)).get(
+                                Runtime.construct_running_task(proc))[1].get(
                                     'port', 8000)))
                 except IndexError:
                     pass
@@ -187,6 +188,7 @@ def get_all_ports():
     def refresh_tasks(running_task=None):
         log_file = running_task if not running_task or 'pid:' not in running_task else None
         process_name = 'swift'
+        negative_name = 'swift.exe'
         cmd_name = 'deploy'
         process = []
         selected = None
@@ -196,9 +198,12 @@ def refresh_tasks(running_task=None):
             except (psutil.ZombieProcess, psutil.AccessDenied,
                     psutil.NoSuchProcess):
                 cmdlines = []
-            if any([process_name in cmdline
-                    for cmdline in cmdlines]) and any(  # noqa
-                        [cmd_name == cmdline for cmdline in cmdlines]):  # noqa
+            if any([
+                    process_name in cmdline for cmdline in cmdlines
+            ]) and not any([negative_name in cmdline
+                            for cmdline in cmdlines]) and any(  # noqa
+                                [cmd_name == cmdline
+                                 for cmdline in cmdlines]):  # noqa
                 process.append(Runtime.construct_running_task(proc))
                 if log_file is not None and any(  # noqa
                     [log_file == cmdline for cmdline in cmdlines]):  # noqa
@@ -240,8 +245,11 @@ def format_time(seconds):
 
     @staticmethod
     def parse_info_from_cmdline(task):
+        pid = None
         for i in range(3):
             slash = task.find('/')
+            if i == 0:
+                pid = task[:slash].split(':')[1]
             task = task[slash + 1:]
         args = task.split('swift deploy')[1]
         args = [arg.strip() for arg in args.split('--') if arg.strip()]
@@ -250,20 +258,23 @@ def parse_info_from_cmdline(task):
             space = args[i].find(' ')
             splits = args[i][:space], args[i][space + 1:]
             all_args[splits[0]] = splits[1]
-        return all_args
+        return pid, all_args
 
     @staticmethod
     def kill_task(task):
-        all_args = Runtime.parse_info_from_cmdline(task)
+        pid, all_args = Runtime.parse_info_from_cmdline(task)
         log_file = all_args['log_file']
-        os.system(f'pkill -9 -f {log_file}')
+        if sys.platform == 'win32':
+            os.system(f'taskkill /f /t /pid "{pid}"')
+        else:
+            os.system(f'pkill -9 -f {log_file}')
         time.sleep(1)
         return [Runtime.refresh_tasks()] + [gr.update(value=None)]
 
     @staticmethod
     def task_changed(task, base_tab):
         if task:
-            all_args = Runtime.parse_info_from_cmdline(task)
+            _, all_args = Runtime.parse_info_from_cmdline(task)
         else:
             all_args = {}
         elements = [

swift/ui/llm_train/runtime.py

@@ -1,5 +1,6 @@
 import collections
 import os.path
+import sys
 import time
 import webbrowser
 from datetime import datetime
@@ -330,6 +331,7 @@ def close_tb(logging_dir):
     def refresh_tasks(running_task=None):
         output_dir = running_task if not running_task or 'pid:' not in running_task else None
         process_name = 'swift'
+        negative_name = 'swift.exe'
         cmd_name = 'sft'
         process = []
         selected = None
@@ -339,9 +341,12 @@ def refresh_tasks(running_task=None):
             except (psutil.ZombieProcess, psutil.AccessDenied,
                     psutil.NoSuchProcess):
                 cmdlines = []
-            if any([process_name in cmdline
-                    for cmdline in cmdlines]) and any(  # noqa
-                        [cmd_name == cmdline for cmdline in cmdlines]):  # noqa
+            if any([
+                    process_name in cmdline for cmdline in cmdlines
+            ]) and not any([negative_name in cmdline
+                            for cmdline in cmdlines]) and any(  # noqa
+                                [cmd_name == cmdline
+                                 for cmdline in cmdlines]):  # noqa
                 process.append(Runtime.construct_running_task(proc))
                 if output_dir is not None and any(  # noqa
                     [output_dir == cmdline for cmdline in cmdlines]):  # noqa
@@ -383,8 +388,11 @@ def format_time(seconds):
 
     @staticmethod
     def parse_info_from_cmdline(task):
+        pid = None
         for i in range(3):
             slash = task.find('/')
+            if i == 0:
+                pid = task[:slash].split(':')[1]
             task = task[slash + 1:]
         args = task.split('swift sft')[1]
         args = [arg.strip() for arg in args.split('--') if arg.strip()]
@@ -406,13 +414,16 @@ def parse_info_from_cmdline(task):
                             f'"{value}"' for value in all_args[key]
                         ]
                     all_args[key] = ' '.join(all_args[key])
-        return all_args
+        return pid, all_args
 
     @staticmethod
     def kill_task(task):
-        all_args = Runtime.parse_info_from_cmdline(task)
+        pid, all_args = Runtime.parse_info_from_cmdline(task)
         output_dir = all_args['output_dir']
-        os.system(f'pkill -9 -f {output_dir}')
+        if sys.platform == 'win32':
+            os.system(f'taskkill /f /t /pid "{pid}"')
+        else:
+            os.system(f'pkill -9 -f {output_dir}')
         time.sleep(1)
         return [Runtime.refresh_tasks()] + [gr.update(value=None)] * (
             len(Runtime.sft_plot) + 1)
@@ -424,7 +435,7 @@ def reset():
     @staticmethod
     def task_changed(task, base_tab):
         if task:
-            all_args = Runtime.parse_info_from_cmdline(task)
+            _, all_args = Runtime.parse_info_from_cmdline(task)
         else:
             all_args = {}
         elements = [
@@ -447,7 +458,7 @@ def task_changed(task, base_tab):
     def plot(task):
         if not task:
             return [None] * len(Runtime.sft_plot)
-        all_args = Runtime.parse_info_from_cmdline(task)
+        _, all_args = Runtime.parse_info_from_cmdline(task)
         tb_dir = all_args['logging_dir']
         fname = [
             fname for fname in os.listdir(tb_dir)