testing a runtime build for apex

Dockerfile

@@ -39,14 +39,7 @@ ENV PYTHONPATH=$HOME/app \
 RUN echo "Installing requirements.txt"
 RUN pip3 install --no-cache-dir --upgrade -r /code/requirements.txt
 
-# Install NVIDIA Apex with CUDA and C++ extensions
-RUN cd $HOME && \
-    git clone https://github.com/NVIDIA/apex && \
-    cd apex && \
-    # sadly this command works at compile time and can be installed, but at runtime it doesn't work (No module named 'fused_layer_norm_cuda')
-    #NVCC_APPEND_FLAGS="--threads 4" pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --global-option="--cpp_ext" --global-option="--cuda_ext" --global-option="--parallel" --global-option="8" ./
-    # so let's try a python only build:
-    pip install -v --disable-pip-version-check --no-build-isolation --no-cache-dir ./
+# Apex will be installed at runtime in run_hf_space.py to ensure CUDA compatibility
 
 WORKDIR $HOME/app
 

run_hf_space.py

@@ -21,6 +21,52 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
+def install_apex():
+    """Install NVIDIA Apex at runtime with CUDA support"""
+    try:
+        logger.info("Installing NVIDIA Apex...")
+        
+        # Clone the Apex repository
+        subprocess.check_call([
+            "git", "clone", "https://github.com/NVIDIA/apex"
+        ])
+        
+        # Change to apex directory and install
+        os.chdir("apex")
+        
+        # Try to install with CUDA extensions first
+        try:
+            logger.info("Attempting to install Apex with CUDA extensions...")
+            subprocess.check_call([
+                sys.executable, "-m", "pip", "install", "-v", 
+                "--disable-pip-version-check", "--no-cache-dir", 
+                "--no-build-isolation", "--global-option=--cpp_ext", 
+                "--global-option=--cuda_ext", "./"
+            ])
+            logger.info("Apex installed successfully with CUDA extensions!")
+        except subprocess.CalledProcessError as e:
+            logger.warning(f"Failed to install Apex with CUDA extensions: {e}")
+            logger.info("Falling back to Python-only build...")
+            
+            # Fall back to Python-only build
+            subprocess.check_call([
+                sys.executable, "-m", "pip", "install", "-v",
+                "--disable-pip-version-check", "--no-build-isolation",
+                "--no-cache-dir", "./"
+            ])
+            logger.info("Apex installed successfully (Python-only build)!")
+            
+    except subprocess.CalledProcessError as e:
+        logger.error(f"Failed to install Apex. Error: {e}")
+        # Don't fail the entire startup if Apex installation fails
+        logger.warning("Continuing without Apex...")
+    except Exception as e:
+        logger.error(f"Unexpected error during Apex installation: {e}")
+        logger.warning("Continuing without Apex...")
+    finally:
+        # Change back to original directory
+        os.chdir("..")
+
 async def run_async():
     """Run the server using the async API directly for better stability"""
     # Get the port from environment variable in Hugging Face Space
@@ -89,6 +135,9 @@ def main():
     subprocess.run(cmd, shell=True)
 
 if __name__ == "__main__":
+    # Install Apex at runtime before starting the server
+    install_apex()
+    
     # First try to run using the async API
     try:
         logger.info("Starting server using async API")