From 1038bb39a1a23d592cc066b7f1b2607c20528f94 Mon Sep 17 00:00:00 2001
From: Meng Zhang <wsxiaoys.lh@gmail.com>
Date: Sat, 25 Mar 2023 00:05:47 +0800
Subject: [PATCH] Add dagu for data processing job orchestration (#7)

* Install dagu

* Move dagu install to first stage

* Fix metrics

* Add DAGs for create dataset from code repository
---
 .dockerignore                                 |  1 +
 Makefile                                      |  2 +-
 README.md                                     |  1 -
 deployment/Dockerfile                         |  8 +++++
 deployment/config/dags/train.yaml             | 20 +++++++++++
 deployment/docker-compose.dev.yml             |  6 ++++
 deployment/docker-compose.yml                 |  9 +++++
 scripts/download_project.py                   | 35 +++++++++++++++++++
 .../huggingface_gptneox_convert.py            |  0
 {tasks => scripts}/preprocess/args.py         |  0
 {tasks => scripts}/preprocess/filters.py      |  0
 {tasks => scripts}/preprocess/metrics.py      | 12 +++++--
 .../preprocess/preprocess_project.py          |  6 +++-
 ...gramming-languages-to-file-extensions.json |  0
 tasks/README.md                               |  5 ---
 15 files changed, 94 insertions(+), 11 deletions(-)
 create mode 100644 deployment/config/dags/train.yaml
 create mode 100644 scripts/download_project.py
 rename {tasks/converter => scripts}/huggingface_gptneox_convert.py (100%)
 rename {tasks => scripts}/preprocess/args.py (100%)
 rename {tasks => scripts}/preprocess/filters.py (100%)
 rename {tasks => scripts}/preprocess/metrics.py (64%)
 rename {tasks => scripts}/preprocess/preprocess_project.py (90%)
 rename {tasks => scripts}/preprocess/programming-languages-to-file-extensions.json (100%)
 delete mode 100644 tasks/README.md

diff --git a/.dockerignore b/.dockerignore
index 8a66279..d0b1763 100644
--- a/.dockerignore
+++ b/.dockerignore
@@ -1,4 +1,5 @@
 testdata
+deployment
 
 **/.git
 **/__pycache__
diff --git a/Makefile b/Makefile
index 1b1968b..d849eb2 100644
--- a/Makefile
+++ b/Makefile
@@ -15,7 +15,7 @@ $(PRE_COMMIT_HOOK):
 	poetry run pre-commit install --install-hooks
 
 $(LOCAL_MODEL):
-	poetry run python tasks/converter/huggingface_gptneox_convert.py \
+	poetry run python scripts/huggingface_gptneox_convert.py \
 		-in_file EleutherAI/pythia-70m-deduped \
 		-o $@ \
 		-i_g 1 -m_n tiny-70M -p 1 -w fp16
diff --git a/README.md b/README.md
index 9869647..d5a083b 100644
--- a/README.md
+++ b/README.md
@@ -12,7 +12,6 @@ An opensource / on-prem alternative to GitHub Copilot
 * [`admin`](./admin): Admin panel for monitoring / settings purpose.
 * [`server`](./server): API server for completion requests. It also logs users' selections (as feedback to model's quality).
 * [`deployment`](./deployment): Container related deployment configs.
-* [`tasks`](./tasks): Various data processing scripts.
 
 ## Development
 
diff --git a/deployment/Dockerfile b/deployment/Dockerfile
index 08b06d2..c0412b0 100644
--- a/deployment/Dockerfile
+++ b/deployment/Dockerfile
@@ -2,6 +2,14 @@
 
 FROM pytorch/pytorch:2.0.0-cuda11.7-cudnn8-runtime
 
+# Install dagu (https://github.com/yohamta/dagu)
+RUN <<EOF
+  python -c "import urllib.request; urllib.request.urlretrieve('https://github.com/yohamta/dagu/releases/download/v1.10.2/dagu_1.10.2_Linux_x86_64.tar.gz', 'dagu.tar.gz')"
+  tar zxvf dagu.tar.gz
+  mv dagu /usr/local/bin
+  rm dagu.tar.gz LICENSE.md README.md
+EOF
+
 ARG PYPI_INDEX_URL=https://pypi.org/simple
 ARG POETRY_VERSION=1.4.0
 
diff --git a/deployment/config/dags/train.yaml b/deployment/config/dags/train.yaml
new file mode 100644
index 0000000..cd66e8c
--- /dev/null
+++ b/deployment/config/dags/train.yaml
@@ -0,0 +1,20 @@
+schedule: "5 4 * * *" # Run at 04:05.
+
+env:
+  - APP_DIR: /app
+steps:
+  - name: Create job dir
+    command: mktemp -d /tmp/job.XXXXX
+    output: JOB_DIR
+
+  - name: Fetch source code
+    dir: $APP_DIR
+    command: python scripts/download_project.py --url http://download.savannah.gnu.org/releases/tinycc/tcc-0.9.24.tar.bz2 --output_dir $JOB_DIR/repository
+    depends:
+      - Create job dir
+
+  - name: Generate dataset
+    dir: ${APP_DIR}
+    command: python scripts/preprocess/preprocess_project.py --project_dir $JOB_DIR/repository --output_dir $JOB_DIR/dataset
+    depends:
+      - Fetch source code
diff --git a/deployment/docker-compose.dev.yml b/deployment/docker-compose.dev.yml
index 052848e..9a2e5d2 100644
--- a/deployment/docker-compose.dev.yml
+++ b/deployment/docker-compose.dev.yml
@@ -19,3 +19,9 @@ services:
   vector:
     environment:
       - VECTOR_WATCH_CONFIG=true
+
+  dagu:
+    ports:
+      - 8080:8080
+    volumes:
+      - ../:/app
diff --git a/deployment/docker-compose.yml b/deployment/docker-compose.yml
index 4fdb86e..83092c9 100644
--- a/deployment/docker-compose.yml
+++ b/deployment/docker-compose.yml
@@ -34,3 +34,12 @@ services:
     volumes:
       - ./config/vector.toml:/etc/vector/vector.toml:ro
       - ./data/logs:/logs
+
+  dagu:
+    image: tabbyml/tabby
+    container_name: tabby-dagu
+    command: dagu server --host 0.0.0.0 --port 8080
+    ports:
+      - 8080:8080
+    volumes:
+      - ./config/dags:/root/.dagu/dags
diff --git a/scripts/download_project.py b/scripts/download_project.py
new file mode 100644
index 0000000..6460bfe
--- /dev/null
+++ b/scripts/download_project.py
@@ -0,0 +1,35 @@
+import tarfile
+from dataclasses import dataclass, field
+from typing import Optional
+
+import requests
+from transformers import HfArgumentParser
+
+
+@dataclass(kw_only=True)
+class DownloaderArgs:
+    url: str = field(metadata={"help": "URL to source code tar.gz file"})
+    output_dir: str = field(metadata={"help": "Output save path directory"})
+
+
+def parse_args():
+    parser = HfArgumentParser(DownloaderArgs)
+    return parser.parse_args()
+
+
+def download_and_untar(url, output_dir):
+    response = requests.get(url, stream=True)
+    mode = "r"
+    if url.endswith(".gz"):
+        mode += "|gz"
+    elif url.endswith(".xz"):
+        mode += "|xz"
+    elif url.endswith(".bz2"):
+        mode += "|bz2"
+    file = tarfile.open(fileobj=response.raw, mode=mode)
+    file.extractall(output_dir)
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    download_and_untar(args.url, args.output_dir)
diff --git a/tasks/converter/huggingface_gptneox_convert.py b/scripts/huggingface_gptneox_convert.py
similarity index 100%
rename from tasks/converter/huggingface_gptneox_convert.py
rename to scripts/huggingface_gptneox_convert.py
diff --git a/tasks/preprocess/args.py b/scripts/preprocess/args.py
similarity index 100%
rename from tasks/preprocess/args.py
rename to scripts/preprocess/args.py
diff --git a/tasks/preprocess/filters.py b/scripts/preprocess/filters.py
similarity index 100%
rename from tasks/preprocess/filters.py
rename to scripts/preprocess/filters.py
diff --git a/tasks/preprocess/metrics.py b/scripts/preprocess/metrics.py
similarity index 64%
rename from tasks/preprocess/metrics.py
rename to scripts/preprocess/metrics.py
index 6a062fd..8fd82ab 100644
--- a/tasks/preprocess/metrics.py
+++ b/scripts/preprocess/metrics.py
@@ -1,16 +1,22 @@
 def max_line_length(content):
-    return max([len(x) for x in content.splitlines()])
+    return max([0] + [len(x) for x in content.splitlines()])
 
 
 def avg_line_length(content):
     lines = [len(x) for x in content.splitlines()]
     total = sum(lines)
-    return total / len(lines)
+    if len(lines) != 0:
+        return total / len(lines)
+    else:
+        return 0
 
 
 def alphanum_fraction(content):
     alphanum = [x for x in content if x.isalpha() or x.isnumeric()]
-    return len(alphanum) / len(content)
+    if len(content) != 0:
+        return len(alphanum) / len(content)
+    else:
+        return 0
 
 
 def compute(content):
diff --git a/tasks/preprocess/preprocess_project.py b/scripts/preprocess/preprocess_project.py
similarity index 90%
rename from tasks/preprocess/preprocess_project.py
rename to scripts/preprocess/preprocess_project.py
index d9c6580..7ec3bc2 100644
--- a/tasks/preprocess/preprocess_project.py
+++ b/scripts/preprocess/preprocess_project.py
@@ -45,7 +45,11 @@ def dataset_iter(files):
             _, extname = os.path.splitext(x)
 
             with open(x) as f:
-                content = f.read()
+                try:
+                    content = f.read()
+                except UnicodeDecodeError:
+                    print("Cannot decode unicode", x)
+                    continue
 
             yield dict(
                 language=mappings[extname],
diff --git a/tasks/preprocess/programming-languages-to-file-extensions.json b/scripts/preprocess/programming-languages-to-file-extensions.json
similarity index 100%
rename from tasks/preprocess/programming-languages-to-file-extensions.json
rename to scripts/preprocess/programming-languages-to-file-extensions.json
diff --git a/tasks/README.md b/tasks/README.md
deleted file mode 100644
index 89cf0d0..0000000
--- a/tasks/README.md
+++ /dev/null
@@ -1,5 +0,0 @@
-## Tasks
-
-* [`converter`](./converter): Converts a [transformers](https://huggingface.co/docs/transformers) causal LM model into TensorRT / FasterTransformer serving formats.
-* [`preprocess`](./preprocess): Preprocess files into [datasets](https://huggingface.co/docs/datasets)
-* [`tabformer`](./tabformer): *NOT RELEASED* Trainer(PEFT w/RLHF) for tabby models.