From c7f13438c9ddbb9f95a1a767ad8458ee47709c14 Mon Sep 17 00:00:00 2001
From: S1ro1 <matej.sirovatka@gmail.com>
Date: Sat, 9 Nov 2024 13:10:18 +0100
Subject: [PATCH] Fix: stuff

---
 .github/workflows/train_workflow.yml | 17 ++++++---------
 discord-bot.py                       | 31 ++++++++++++++++------------
 2 files changed, 24 insertions(+), 24 deletions(-)

diff --git a/.github/workflows/train_workflow.yml b/.github/workflows/train_workflow.yml
index 674aee64..d3a74151 100644
--- a/.github/workflows/train_workflow.yml
+++ b/.github/workflows/train_workflow.yml
@@ -18,11 +18,6 @@ jobs:
   train:
     runs-on: ubuntu-latest
     steps:
-      - name: Set up environment variables
-        run: |
-          echo "SCRIPT_FILE=train.${{ inputs.script_type }}" >> $GITHUB_ENV
-          echo "OUTPUT_FILE=training.log" >> $GITHUB_ENV
-
       - name: Install Python dependencies
         if: inputs.script_type == 'py'
         run: |
@@ -38,23 +33,23 @@ jobs:
 
       - name: Create training script
         run: |
-          echo "${{ inputs.script_content }}" > ${{ env.SCRIPT_FILE }}
-          cat ${{ env.SCRIPT_FILE }}  # Debug: print the content
+          echo '${{ inputs.script_content }}' > train.${{ inputs.script_type }}
+          cat train.${{ inputs.script_type }}
 
       - name: Compile and run CUDA script
         if: inputs.script_type == 'cu'
         run: |
-          nvcc ${{ env.SCRIPT_FILE }} -o train_cuda
-          ./train_cuda > ${{ env.OUTPUT_FILE }} 2>&1
+          nvcc train.cu -o train_cuda
+          ./train_cuda > training.log 2>&1
 
       - name: Run Python script
         if: inputs.script_type == 'py'
         run: |
-          python ${{ env.SCRIPT_FILE }} > ${{ env.OUTPUT_FILE }} 2>&1
+          python train.py > training.log 2>&1
 
       - name: Upload logs
         uses: actions/upload-artifact@v3
         if: always()
         with:
           name: training-logs
-          path: ${{ env.OUTPUT_FILE }}
+          path: training.log
diff --git a/discord-bot.py b/discord-bot.py
index 85b0bb57..9ff6ca48 100644
--- a/discord-bot.py
+++ b/discord-bot.py
@@ -59,7 +59,7 @@ def get_github_branch_name():
 client = discord.Client(intents=intents)
 
 
-async def trigger_github_action(script_content):
+async def trigger_github_action(script_content, script_type):
     """
     Triggers the GitHub action with custom train.py contents
     """
@@ -74,7 +74,7 @@ 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, 'script_type': script_type})
         logger.info(f"Workflow dispatch result: {success}")
         
         if success:
@@ -173,25 +173,29 @@ async def on_ready():
             logger.warning(f'Failed to update nickname in guild {guild.name}: {e}')
 
 
-async def process_python(message, attachment):
+async def process_file(message, attachment, script_type):
+    """
+    Generic file processor that handles both .py and .cu files
+    """
     # Reply to the original message
-    initial_reply = await message.reply("Found train.py! Starting training process...")
+    initial_reply = await message.reply(f"Found train.{script_type}! Starting training process...")
     
     # Create a new thread from the reply
     thread = await initial_reply.create_thread(
         name=f"Training Job - {datetime.now().strftime('%Y-%m-%d %H:%M')}",
-        auto_archive_duration=1440  # Archive after 24 hours of inactivity
+        auto_archive_duration=1440
     )
     
     try:
         # Download the file content
-        logger.info("Downloading train.py content")
+        logger.info(f"Downloading train.{script_type} content")
         script_content = await attachment.read()
         script_content = script_content.decode('utf-8')
-        logger.info("Successfully read train.py content")
+        logger.info(script_content)
+        logger.info(f"Successfully read train.{script_type} content")
         
-        # Trigger GitHub Action
-        run_id = await trigger_github_action(script_content)
+        # Trigger GitHub Action with appropriate script_type
+        run_id = await trigger_github_action(script_content, script_type)
         
         if run_id:
             logger.info(f"Successfully triggered workflow with run ID: {run_id}")
@@ -235,13 +239,14 @@ async def on_message(message):
             for attachment in message.attachments:
                 logger.info(f"Processing attachment: {attachment.filename}")
                 if attachment.filename == "train.py":
-                    process_python(message, attachment)
+                    await process_file(message, attachment, 'py')
                     break
                 elif attachment.filename == "train.cu":
-                    raise ValueError("CUDA training is not supported yet")
+                    await process_file(message, attachment, 'cu')
+                    break
 
-            if not any(att.filename == "train.py" for att in message.attachments):
-                await message.reply("Please attach a file named 'train.py' to your message.")
+            if not any(att.filename in ["train.py", "train.cu"] for att in message.attachments):
+                await message.reply("Please attach a file named 'train.py' or 'train.cu' to your message.")
 
 # Run the bot
 if __name__ == "__main__":