subprocess.Run

app.py

@@ -31,6 +31,7 @@ class QuantizationConfig:
     quantized_gguf: str = field(default="", init=False)
     imatrix_file: str = field(default="", init=False)
 
+
 @dataclass
 class SplitConfig:
     """Configuration for model splitting."""
@@ -46,6 +47,7 @@ class OutputConfig:
     repo_name: str = ""
     filename: str = ""
 
+
 @dataclass
 class ModelProcessingConfig:
     """Configuration for the entire model processing pipeline."""
@@ -60,6 +62,7 @@ class ModelProcessingConfig:
     new_repo_url: str = field(default="", init=False)
     new_repo_id: str = field(default="", init=False)
 
+
 class GGUFConverterError(Exception):
     """Custom exception for GGUF conversion errors."""
     pass
@@ -143,9 +146,10 @@ class HuggingFaceModelProcessor:
         train_data_path = self.CALIBRATION_FILE
         if not os.path.isfile(train_data_path):
             raise GGUFConverterError(f"Training data file not found: {train_data_path}")
-        print(f"Training data file path: {train_data_path}")
 
+        print(f"Training data file path: {train_data_path}")
         print("Running imatrix command...")
+
         imatrix_command = [
             "llama-imatrix",
             "-m", quant_config.fp16_model,
@@ -157,16 +161,19 @@ class HuggingFaceModelProcessor:
 
         process = subprocess.Popen(imatrix_command, shell=False)
         try:
-            process.wait(timeout=60)
+            process.wait(timeout=300)
         except subprocess.TimeoutExpired:
             print("Imatrix computation timed out. Sending SIGINT to allow graceful termination...")
             process.send_signal(signal.SIGINT)
             try:
                 process.wait(timeout=5)
             except subprocess.TimeoutExpired:
-                print("Imatrix proc still didn't term. Forecfully terming process...")
+                print("Imatrix proc still didn't term. Forcefully terminating process...")
                 process.kill()
 
+        if process.returncode != 0:
+             raise GGUFConverterError(f"Error generating imatrix")
+
         print(f"Importance matrix generation completed: {os.path.abspath(quant_config.imatrix_file)}")
 
     def _split_and_upload_model(self, processing_config: ModelProcessingConfig) -> None:
@@ -188,14 +195,16 @@ class HuggingFaceModelProcessor:
         split_cmd.extend([quant_config.quantized_gguf, model_path_prefix])
 
         print(f"Split command: {split_cmd}")
-        result = subprocess.run(split_cmd, shell=False, capture_output=True, text=True)
-
-        print(f"Split command stdout: {result.stdout}")
-        print(f"Split command stderr: {result.stderr}")
+        process = subprocess.Popen(split_cmd, shell=False)
+        try:
+            process.wait(timeout=300)
+        except subprocess.TimeoutExpired:
+            print("Splitting timed out. Killing process...")
+            process.kill()
+            raise GGUFConverterError("Error splitting the model: Operation timed out.")
 
-        if result.returncode != 0:
-            stderr_str = result.stderr.decode("utf-8")
-            raise GGUFConverterError(f"Error splitting the model: {stderr_str}")
+        if process.returncode != 0:
+            raise GGUFConverterError(f"Error splitting the model")
 
         print("Model split successfully!")
 
@@ -215,6 +224,7 @@ class HuggingFaceModelProcessor:
             raise GGUFConverterError("No sharded files found.")
 
         print(f"Sharded model files: {sharded_model_files}")
+
         for file in sharded_model_files:
             file_path = os.path.join(processing_config.outdir, file)
             try:
@@ -268,19 +278,20 @@ class HuggingFaceModelProcessor:
 
             # Convert HF to GGUF
             print(f"Converting to GGUF FP16: {os.path.abspath(processing_config.quant_config.fp16_model)}")
-            result = subprocess.run(
-                [
-                    "python3", "/app/convert_hf_to_gguf.py", local_dir,
-                    "--outtype", "f16", "--outfile", processing_config.quant_config.fp16_model
-                ],
-                shell=False,
-                capture_output=True
-            )
+            convert_command = [
+                "python3", "/app/convert_hf_to_gguf.py", local_dir,
+                "--outtype", "f16", "--outfile", processing_config.quant_config.fp16_model
+            ]
+            process = subprocess.Popen(convert_command, shell=False)
+            try:
+                process.wait(timeout=600)
+            except subprocess.TimeoutExpired:
+                print("Conversion timed out. Killing process...")
+                process.kill()
+                raise GGUFConverterError("Error converting to fp16: Operation timed out.")
 
-            print(f"Model directory contents: {result}")
-            if result.returncode != 0:
-                stderr_str = result.stderr.decode("utf-8")
-                raise GGUFConverterError(f"Error converting to fp16: {stderr_str}")
+            if process.returncode != 0:
+                raise GGUFConverterError(f"Error converting to fp16")
 
             print("Model converted to fp16 successfully!")
             print(f"Converted model path: {os.path.abspath(processing_config.quant_config.fp16_model)}")
@@ -315,11 +326,18 @@ class HuggingFaceModelProcessor:
             quantize_cmd.append(quant_config.method)
 
         print(f"Quantizing model with {quantize_cmd}")
-        result = subprocess.run(quantize_cmd, shell=False, capture_output=True)
 
-        if result.returncode != 0:
-            stderr_str = result.stderr.decode("utf-8")
-            raise GGUFConverterError(f"Error quantizing: {stderr_str}")
+        # Use Popen for quantization
+        process = subprocess.Popen(quantize_cmd, shell=False)
+        try:
+            process.wait(timeout=3600)
+        except subprocess.TimeoutExpired:
+            print("Quantization timed out. Killing process...")
+            process.kill()
+            raise GGUFConverterError("Error quantizing: Operation timed out.")
+
+        if process.returncode != 0:
+            raise GGUFConverterError(f"Error quantizing")
 
         print(f"Quantized successfully with {quant_config.imatrix_method if quant_config.use_imatrix else quant_config.method} option!")
         print(f"Quantized model path: {os.path.abspath(quant_config.quantized_gguf)}")
@@ -338,7 +356,7 @@ class HuggingFaceModelProcessor:
 
         return new_repo_url
 
-    def _generate_readme(self, processing_config :ModelProcessingConfig) -> str:
+    def _generate_readme(self, processing_config: ModelProcessingConfig) -> str:
         """Generate README.md for the quantized model."""
         creator = self._get_model_creator(processing_config.model_id)
         username = whoami(processing_config.token)["name"]