Generalize file name

.github/workflows/train_workflow.yml

@@ -4,7 +4,11 @@ on:
   workflow_dispatch:
     inputs:
       script_content:
-        description: 'Content of train.py'
+        description: 'Content of Python script'
+        required: true
+        type: string
+      filename:
+        description: 'Name of Python script'
         required: true
         type: string
 
@@ -19,32 +23,32 @@ jobs:
       with:
         python-version: '3.10'
 
-    - name: Create training script
+    - name: Create script
       shell: python
       run: |
-        with open('train.py', 'w') as f:
+        with open('${{ github.event.inputs.filename }}', 'w') as f:
             f.write('''${{ github.event.inputs.script_content }}''')
 
     - name: Install dependencies
       run: |
-        if grep -rE "(import numpy|from numpy)" train.py; then
+        if grep -rE "(import numpy|from numpy)" "${{ github.event.inputs.filename }}"; then
           echo "Numpy detected, installing numpy"
           pip install numpy
         fi
         # Check if 'import torch' is in any Python file
-        if grep -rE "(import torch|from torch)" train.py; then
+        if grep -rE "(import torch|from torch)" "${{ github.event.inputs.filename }}"; then
           echo "PyTorch detected, installing torch"
           pip install torch
         fi
         # Check if 'import triton' is in any Python file
-        if grep -rE "(import triton|from triton)" train.py; then
+        if grep -rE "(import triton|from triton)" "${{ github.event.inputs.filename }}"; then
           echo "Triton detected, installing triton"
           pip install triton
         fi
 
-    - name: Run training script
+    - name: Run script
       run: |
-        python train.py > training.log 2>&1
+        python "${{ github.event.inputs.filename }}" > training.log 2>&1
 
     - name: Upload artifacts
       uses: actions/upload-artifact@v3
@@ -53,6 +57,6 @@ jobs:
         name: training-artifacts
         path: |
           training.log
-          train.py
+          ${{ github.event.inputs.filename }}
       env:
         CUDA_VISIBLE_DEVICES: 0 # Make sure only one GPU is used for testing

.gitignore

@@ -1,3 +1,4 @@
 .env
 __pycache__/
-*.pyc
+*.pyc
+.DS_Store

discord-bot.py

@@ -59,9 +59,9 @@ def get_github_branch_name():
 client = discord.Client(intents=intents)
 
 
-async def trigger_github_action(script_content):
+async def trigger_github_action(script_content, filename):
     """
-    Triggers the GitHub action with custom train.py contents
+    Triggers the GitHub action with custom script contents and filename
     """
     logger.info("Attempting to trigger GitHub action")
     gh = Github(os.getenv('GITHUB_TOKEN'))
@@ -74,7 +74,10 @@ async def trigger_github_action(script_content):
         workflow = repo.get_workflow("train_workflow.yml")
         logger.info("Found workflow, attempting to dispatch")
 
-        success = workflow.create_dispatch(get_github_branch_name(), {'script_content': script_content})
+        success = workflow.create_dispatch(get_github_branch_name(), {
+            'script_content': script_content,
+            'filename': filename
+        })
         logger.info(f"Workflow dispatch result: {success}")
 
         if success:
@@ -195,25 +198,25 @@ async def on_message(message):
         if message.attachments:
             for attachment in message.attachments:
                 logger.info(f"Processing attachment: {attachment.filename}")
-                if attachment.filename == "train.py":
+                if attachment.filename.endswith('.py'):
                     # Create a thread directly from the original message
                     thread = await message.create_thread(
                         name=f"Training Job - {datetime.now().strftime('%Y-%m-%d %H:%M')}",
                         auto_archive_duration=1440  # Archive after 24 hours of inactivity
                     )
 
                     # Send initial message in the thread
-                    await thread.send("Found train.py! Starting training process...")
+                    await thread.send(f"Found {attachment.filename}! Starting training process...")
 
                     try:
                         # Download the file content
-                        logger.info("Downloading train.py content")
+                        logger.info(f"Downloading {attachment.filename} content")
                         script_content = await attachment.read()
                         script_content = script_content.decode('utf-8')
-                        logger.info("Successfully read train.py content")
+                        logger.info(f"Successfully read {attachment.filename} content")
 
                         # Trigger GitHub Action
-                        run_id = await trigger_github_action(script_content)
+                        run_id = await trigger_github_action(script_content, attachment.filename)
 
                         if run_id:
                             logger.info(f"Successfully triggered workflow with run ID: {run_id}")
@@ -236,7 +239,7 @@ async def on_message(message):
                             if url:
                                 await thread.send(f"View the full run at: {url}")
                         else:
-                            logger.error("Failed to trigger GitHub Action")
+                            logger.error("Missing run_id. Failed to trigger GitHub Action")
                             await thread.send("Failed to trigger GitHub Action. Please check the configuration.")
 
                     except Exception as e:
@@ -245,7 +248,7 @@ async def on_message(message):
 
                     break
 
-            if not any(att.filename == "train.py" for att in message.attachments):
+            if not any(att.filename.endswith('.py') for att in message.attachments):
                 await message.reply("Please attach a file named 'train.py' to your message.")
 
 # Run the bot