160 lines
5.1 KiB
Python
160 lines
5.1 KiB
Python
import os
|
|
import subprocess
|
|
from enum import Enum
|
|
from pathlib import Path
|
|
|
|
|
|
class ModelType(Enum):
|
|
SDXL = "SDXL"
|
|
FLUX = "FLUX"
|
|
StudioVoice = "StudioVoice"
|
|
|
|
class NIMManager:
|
|
'''
|
|
This class is responsible for managing the NIM containers.
|
|
|
|
Supports:
|
|
- SDXL
|
|
- FLUX
|
|
- StudioVoice (https://build.nvidia.com/nvidia/studiovoice/docker)
|
|
'''
|
|
# Registry paths for different model types
|
|
MODEL_REGISTRY: dict[ModelType, str] = {
|
|
ModelType.SDXL: "nvcr.io/nvstaging/nim/sdxl:1.0.0",
|
|
ModelType.FLUX: "nvcr.io/nvstaging/nim/flux:1.0.0",
|
|
ModelType.StudioVoice: "nvcr.io/nim/nvidia/maxine-studio-voice:latest"
|
|
}
|
|
|
|
def __init__(self):
|
|
self._nim_server_proc_dict: dict[str, subprocess.Popen] = {}
|
|
|
|
def setup_directories(self, model_name: str, version: str) -> bool:
|
|
"""Create necessary directories for NIM cache"""
|
|
try:
|
|
home = os.path.expanduser("~")
|
|
cache_path = Path(f"{home}/nimcache/{model_name}/{version}/.cache")
|
|
cache_path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Set permissions
|
|
chmod_command = f"chmod 777 -R {cache_path}"
|
|
result = subprocess.run(chmod_command, shell=True, check=True)
|
|
|
|
print(f"Directory setup completed for {model_name}")
|
|
return True
|
|
|
|
except Exception as e:
|
|
print(f"Failed to setup directories: {str(e)}")
|
|
return False
|
|
|
|
def pull_nim_image(self, nim_id: str, registry_path: str) -> bool:
|
|
"""Pull NIM image from the registry"""
|
|
try:
|
|
command = f"wsl -d NVIDIA-Workbench -- podman pull {registry_path}"
|
|
result = subprocess.run(
|
|
command,
|
|
stdout=subprocess.DEVNULL,
|
|
stderr=subprocess.DEVNULL,
|
|
shell=True
|
|
)
|
|
|
|
if result.returncode == 0:
|
|
print(f"{nim_id} image downloaded successfully")
|
|
return True
|
|
else:
|
|
print(f"Failed to download image {nim_id}")
|
|
return False
|
|
|
|
except Exception as e:
|
|
print(f"Error pulling image: {str(e)}")
|
|
return False
|
|
|
|
def start_nim_container(self,
|
|
nim_id: str,
|
|
model_name: str,
|
|
version: str,
|
|
port: int,
|
|
api_key: str,
|
|
registry_path: str) -> bool:
|
|
"""Start a NIM container with the specified configuration"""
|
|
try:
|
|
home = os.path.expanduser("~")
|
|
cache_path = f"{home}/nimonwsl2/{model_name}/{version}/.cache"
|
|
|
|
command = (
|
|
f"wsl -d NVIDIA-Workbench -- podman run -it --rm "
|
|
f"--device=nvidia.com/gpu=all "
|
|
f"--name={model_name} "
|
|
f"--shm-size=8GB "
|
|
f"-e NGC_API_KEY={api_key} "
|
|
f"-v {cache_path}:/opt/nim/.cache "
|
|
f"-e NIM_RELAX_MEM_CONSTRAINTS=1 "
|
|
f"-p {port}:8000 "
|
|
f"{registry_path}"
|
|
)
|
|
|
|
run_process = subprocess.Popen(
|
|
command,
|
|
stdin=subprocess.PIPE,
|
|
stdout=subprocess.DEVNULL,
|
|
stderr=subprocess.DEVNULL,
|
|
shell=True
|
|
)
|
|
|
|
if run_process:
|
|
self._nim_server_proc_dict[nim_id] = run_process
|
|
print(f"Started NIM container for {nim_id} on port {port}")
|
|
return True
|
|
else:
|
|
print(f"Failed to start container for {nim_id}")
|
|
return False
|
|
|
|
except Exception as e:
|
|
print(f"Error starting container: {str(e)}")
|
|
return False
|
|
|
|
def deploy_nim(self,
|
|
model_name: ModelType,
|
|
version: str,
|
|
port: int,
|
|
api_key: str) -> bool:
|
|
"""Deploy a NIM model with all necessary setup"""
|
|
nim_id = f"{model_name.value}-{version}"
|
|
|
|
# Get registry path from MODEL_REGISTRY
|
|
registry_path = self.MODEL_REGISTRY[model_name]
|
|
|
|
# Setup directories
|
|
if not self.setup_directories(model_name.value, version):
|
|
return False
|
|
|
|
# Pull image
|
|
if not self.pull_nim_image(nim_id, registry_path):
|
|
return False
|
|
|
|
# Start container
|
|
return self.start_nim_container(
|
|
nim_id=nim_id,
|
|
model_name=model_name.value,
|
|
version=version,
|
|
port=port,
|
|
api_key=api_key,
|
|
registry_path=registry_path
|
|
)
|
|
|
|
# Example usage
|
|
if __name__ == "__main__":
|
|
manager = NIMManager()
|
|
|
|
# Example deployment
|
|
success = manager.deploy_nim(
|
|
model_name=ModelType.SDXL,
|
|
version="1.5.1-rtx.rc3",
|
|
port=18001,
|
|
api_key="your-api-key-here"
|
|
)
|
|
|
|
if success:
|
|
print("NIM deployment successful")
|
|
else:
|
|
print("NIM deployment failed")
|