| import subprocess |
|
|
| import typer |
|
|
| import docker |
|
|
| docker_client = docker.from_env() |
|
|
|
|
| def ask_user_to_kill_container(c): |
| name = c.attrs.get("Name", "") |
| env_vars = c.attrs.get("Config", {}).get("Env", {}) |
| port = 9001 |
| for var in env_vars: |
| if var.startswith("PORT="): |
| port = var.split("=")[1] |
| should_delete = typer.confirm( |
| f" An inference server is already running in container {name} on port {port}. Are you sure you want to delete it?" |
| ) |
| return should_delete |
|
|
|
|
| def is_inference_server_container(container): |
| image_tags = container.image.tags |
| for t in image_tags: |
| if t.startswith("roboflow/roboflow-inference-server"): |
| return True |
| return False |
|
|
|
|
| def handle_existing_containers(containers): |
| has_existing_containers = False |
| for c in containers: |
| if is_inference_server_container(c): |
| has_existing_containers = True |
| if c.attrs.get("State", {}).get("Status", "").lower() == "running": |
| should_kill = ask_user_to_kill_container(c) |
| if should_kill: |
| c.kill() |
| has_existing_containers = False |
| return has_existing_containers |
|
|
|
|
| def find_existing_containers(): |
| containers = [] |
| for c in docker_client.containers.list(): |
| if is_inference_server_container(c): |
| if c.attrs.get("State", {}).get("Status", "").lower() == "running": |
| containers.append(c) |
| return containers |
|
|
|
|
| def get_image(): |
| try: |
| subprocess.check_output("nvidia-smi") |
| print("GPU detected. Using a GPU image.") |
| return "roboflow/roboflow-inference-server-gpu:latest" |
| except: |
| print("No GPU detected. Using a CPU image.") |
| return "roboflow/roboflow-inference-server-cpu:latest" |
|
|
|
|
| def start_inference_container( |
| api_key, |
| image=None, |
| port=9001, |
| labels=None, |
| project="roboflow-platform", |
| metrics_enabled=True, |
| device_id=None, |
| num_workers=1, |
| ): |
| containers = find_existing_containers() |
| if len(containers) > 0: |
| still_has_containers = handle_existing_containers(containers) |
| if still_has_containers: |
| print("Please kill the existing containers and try again.") |
| return |
|
|
| if image is None: |
| image = get_image() |
|
|
| device_requests = None |
| privileged = False |
| if "gpu" in image: |
| privileged = True |
| device_requests = ( |
| [docker.types.DeviceRequest(device_ids=["all"], capabilities=[["gpu"]])], |
| ) |
|
|
| print(f"Starting inference server container...") |
| docker_client.containers.run( |
| image=image, |
| privileged=privileged, |
| detach=True, |
| labels=labels, |
| ports={"9001": port}, |
| |
| device_requests=device_requests, |
| environment=[ |
| "HOST=0.0.0.0", |
| f"PORT={port}", |
| f"PROJECT={project}", |
| f"METRICS_ENABLED={metrics_enabled}", |
| f"DEVICE_ID={device_id}", |
| f"API_KEY={api_key}", |
| f"NUM_WORKERS={num_workers}", |
| ], |
| ) |
|
|
|
|
| def check_inference_server_status(): |
| containers = find_existing_containers() |
| if len(containers) > 0: |
| for c in containers: |
| container_name = c.attrs.get("Name", "") |
| created = c.attrs.get("Created", "") |
| exposed_port = list(c.attrs.get("Config").get("ExposedPorts", {}).keys())[0] |
| status = c.attrs.get("State", {}).get("Status", "unknown") |
| image = c.attrs.get("Image", "") |
| container_status_message = """ |
| Container Name: {container_name} |
| Created: {created} |
| Exposed Port: {exposed_port} |
| Status: {status} |
| Image: {image} |
| """ |
| print( |
| container_status_message.format( |
| container_name=container_name, |
| created=created, |
| exposed_port=exposed_port, |
| status=status, |
| image=image, |
| ) |
| ) |
| return |
| print("No inference server container running.") |
|
|
|
|
| if __name__ == "__main__": |
| start_inference_container("my_api_key") |
|
|
