From dcdd8a5ef7ae450442a20c8021a333501b4922d7 Mon Sep 17 00:00:00 2001
From: Frank Lee <somerlee.9@gmail.com>
Date: Tue, 27 Feb 2024 15:19:13 +0800
Subject: [PATCH] [setup] fixed nightly release (#5388)

---
 .../workflows/release_nightly_on_schedule.yml | 12 ++++---
 .../scripts/update_setup_for_nightly.py       | 34 +++++++++++++++++++
 setup.py                                      | 16 ++-------
 3 files changed, 45 insertions(+), 17 deletions(-)
 create mode 100644 .github/workflows/scripts/update_setup_for_nightly.py

diff --git a/.github/workflows/release_nightly_on_schedule.yml b/.github/workflows/release_nightly_on_schedule.yml
index 4125f333f..072a943ae 100644
--- a/.github/workflows/release_nightly_on_schedule.yml
+++ b/.github/workflows/release_nightly_on_schedule.yml
@@ -6,11 +6,13 @@ on:
     - cron:  '0 0 * * 6' # release on every Sunday 00:00 UTC time
 
 jobs:
-  build-n-publish:
+  publish:
     if: github.repository == 'hpcaitech/ColossalAI'
     name: Build and publish Python 🐍 distributions 📦 to PyPI
     runs-on: ubuntu-latest
     timeout-minutes: 20
+    outputs:
+      status: ${{ steps.publish.outcome }}
     steps:
     - uses: actions/checkout@v2
 
@@ -18,7 +20,9 @@ jobs:
       with:
         python-version: '3.8.14'
 
-    - run: NIGHTLY=1 python setup.py sdist build
+    - run: |
+        python .github/workflows/scripts/update_setup_for_nightly.py
+        python setup.py sdist build
 
     # publish to PyPI if executed on the main branch
     - name: Publish package to PyPI
@@ -31,7 +35,7 @@ jobs:
 
   notify:
     name: Notify Lark via webhook
-    needs: build-n-publish
+    needs: publish
     runs-on: ubuntu-latest
     if: ${{ always() }} && github.repository == 'hpcaitech/ColossalAI'
     steps:
@@ -62,4 +66,4 @@ jobs:
           REPO: ${{ github.repository }}
           RUN_ID: ${{ github.run_id }}
           WEBHOOK_URL: ${{ secrets.LARK_NOTIFICATION_WEBHOOK_URL }}
-          STATUS: ${{ steps.publish.outcome }}
+          STATUS: ${{ needs.publish.outputs.status }}
diff --git a/.github/workflows/scripts/update_setup_for_nightly.py b/.github/workflows/scripts/update_setup_for_nightly.py
new file mode 100644
index 000000000..d8a3087ef
--- /dev/null
+++ b/.github/workflows/scripts/update_setup_for_nightly.py
@@ -0,0 +1,34 @@
+from datetime import datetime
+
+
+def open_setup_file():
+    with open("setup.py", "r") as f:
+        file_lines = f.readlines()
+    return file_lines
+
+
+def replace_nightly_package_info(file_lines):
+    version = datetime.today().strftime("%Y.%m.%d")
+    package_name = "colossalai-nightly"
+
+    for idx, line in enumerate(file_lines):
+        if "version = get_version()" in line:
+            file_lines[idx] = f'version = "{version}"\n'
+        if 'package_name = "colossalai"' in line:
+            file_lines[idx] = f'package_name = "{package_name}"\n'
+    return file_lines
+
+
+def write_setup_file(file_lines):
+    with open("setup.py", "w") as f:
+        f.writelines(file_lines)
+
+
+def main():
+    file_lines = open_setup_file()
+    file_lines = replace_nightly_package_info(file_lines)
+    write_setup_file(file_lines)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/setup.py b/setup.py
index 5f6d447dd..e54ec41ea 100644
--- a/setup.py
+++ b/setup.py
@@ -1,6 +1,5 @@
 import os
 import sys
-from datetime import datetime
 from typing import List
 
 from setuptools import find_packages, setup
@@ -15,7 +14,6 @@ except ImportError:
 
 THIS_DIR = os.path.dirname(os.path.abspath(__file__))
 BUILD_EXT = int(os.environ.get("BUILD_EXT", "0")) == 1
-IS_NIGHTLY = int(os.environ.get("NIGHTLY", "0")) == 1
 
 # we do not support windows currently
 if sys.platform == "win32":
@@ -96,23 +94,15 @@ if BUILD_EXT:
 else:
     ext_modules = []
 
-# always put not nightly branch as the if branch
-# otherwise github will treat colossalai-nightly as the project name
-# and it will mess up with the dependency graph insights
-if not IS_NIGHTLY:
-    version = get_version()
-    package_name = "colossalai"
-else:
-    # use date as the nightly version
-    version = datetime.today().strftime("%Y.%m.%d")
-    package_name = "colossalai-nightly"
+version = get_version()
+package_name = "colossalai"
 
 setup(
     name=package_name,
     version=version,
     packages=find_packages(
         exclude=(
-            "op_builder",
+            "extensions",
             "benchmark",
             "docker",
             "tests",