update cuda path

script.py

@@ -129,19 +129,19 @@ os.environ['MKL_SERVICE_FORCE_INTEL'] = '1'
 # 		print(f"Failed to download packages. Error: {e}")
 
 def download_packages(packages, folder):
-    # Create the directory if it doesn't exist
-    if not os.path.exists(folder):
-        os.makedirs(folder)
-    
-    try:
-        subprocess.check_call([
-            'pip', 'download',
-            '--dest', folder,
-            '-f', 'https://download.pytorch.org/whl/cu121'
-        ] + packages)
-        print(f"Packages downloaded successfully to {folder}")
-    except subprocess.CalledProcessError as e:
-        print(f"Failed to download packages. Error: {e}")
+	# Create the directory if it doesn't exist
+	if not os.path.exists(folder):
+		os.makedirs(folder)
+	
+	try:
+		subprocess.check_call([
+			'pip', 'download',
+			'--dest', folder,
+			'-f', 'https://download.pytorch.org/whl/cu121'
+		] + packages)
+		print(f"Packages downloaded successfully to {folder}")
+	except subprocess.CalledProcessError as e:
+		print(f"Failed to download packages. Error: {e}")
 
 def install_package_from_local_file(package_name, folder='packages'):
 	"""
@@ -170,6 +170,9 @@ def setup_environment():
 	# packages_to_download = ['torch==2.1.0', 'torchvision==0.16.0', 'torchaudio==2.1.0']
 	# download_packages(packages_to_download, folder='packages/torch')
 
+	# Install ninja
+	subprocess.check_call([sys.executable, "-m", "pip", "install", "ninja"])
+
 	# packages_to_download = ['torch==2.1.0', 'torchvision==0.16.0', 'torchaudio==2.1.0']
 	# download_folder = 'packages/torch'
 
@@ -195,6 +198,14 @@ def setup_environment():
 	else:
 		print(f"Directory {pc_util_path} does not exist")
 	
+def setup_cuda_environment():
+	cuda_home = '/usr/local/cuda-12.1'
+	if not os.path.exists(cuda_home):
+		raise EnvironmentError(f"CUDA_HOME directory {cuda_home} does not exist. Please install CUDA and set CUDA_HOME environment variable.")
+	os.environ['CUDA_HOME'] = cuda_home
+	os.environ['PATH'] = f"{cuda_home}/bin:{os.environ['PATH']}"
+	os.environ['LD_LIBRARY_PATH'] = f"{cuda_home}/lib64:{os.environ.get('LD_LIBRARY_PATH', '')}"
+
 from pathlib import Path
 def save_submission(submission, path):
 	"""
@@ -209,6 +220,7 @@ def save_submission(submission, path):
 	print(f"Submission saved to {path}")
 
 if __name__ == "__main__":
+	setup_cuda_environment()
 	setup_environment()
 
 	from handcrafted_solution import predict
@@ -237,4 +249,4 @@ if __name__ == "__main__":
 				# save_submission(solution, Path(params['output_path']) / "submission.parquet")
 	print('------------ Saving results ---------------')
 	save_submission(solution, Path(params['output_path']) / "submission.parquet")
-	print("------------ Done ------------ ")
+	print("------------ Done ------------ ")