WIP

Signed-off-by: troychiu <y.troychiu@gmail.com>

plugins/flytekit-vscode/README.md

@@ -0,0 +1,20 @@
+# Flytekit VSCode Plugin
+
+Flytekit VSCode plugin provides an simple way for users to run the task as a interactive vscode server with any image. `@vscode` is a decorator which users can put within @task and user function. With `@vscode`, the task will install vscode dependencies and run a vscode server instead of the user defined functions.
+
+To install the plugin, run the following command:
+
+```bash
+pip install flytekitplugins-vscode
+```
+
+## Task Example
+```python
+from flytekit import task
+from flytekitplugins.vscode import vscode
+
+@task
+@vscode
+def train():
+    ...
+```

plugins/flytekit-vscode/flytekitplugins/vscode/vscode.py

@@ -1,5 +1,6 @@
 import subprocess
 from functools import wraps
+import sys
 import time
 import os
 import multiprocessing
@@ -130,6 +131,7 @@ def inner_wrapper(*args, **kwargs):
             # 0. Executes the pre_execute function if provided.
             if pre_execute is not None:
                 pre_execute()
+                print_flush("Pre execute function executed successfully!")
 
             # 1. Downloads the VSCode server from Internet to local.
             code_server_bin = download_vscode(
@@ -156,8 +158,7 @@ def inner_wrapper(*args, **kwargs):
                 print_flush("Post execute function executed successfully!")
             child_process.terminate()
             child_process.join()
-
-            return
+            sys.exit(0)
 
         return inner_wrapper