From f4d25cf1911151efaa4b78b6632f8889419e0f9b Mon Sep 17 00:00:00 2001
From: Heiru Wu <heiru.wu@instill.tech>
Date: Sat, 7 Dec 2024 02:21:17 +0800
Subject: [PATCH] feat(cli): allow multi-gpus when running inference locally

---
 instill/helpers/cli.py | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/instill/helpers/cli.py b/instill/helpers/cli.py
index d41d3f9..b1a0483 100644
--- a/instill/helpers/cli.py
+++ b/instill/helpers/cli.py
@@ -118,6 +118,14 @@ def cli():
         action="store_true",
         required=False,
     )
+    run_parser.add_argument(
+        "-ng",
+        "--num-of-gpus",
+        help="number of gpus to use if gpu flag is on, default to 1",
+        type=int,
+        default=1,
+        required=False,
+    )
     run_parser.add_argument(
         "-t",
         "--tag",
@@ -311,7 +319,7 @@ def run(args):
                     {args.name}:{args.tag} /bin/bash -c \
                         \"serve build _model:entrypoint -o serve.yaml && \
                         sed -i 's/app1/default/' serve.yaml && \
-                        sed -i 's/num_cpus: 0.0/num_gpus: 1.0/' serve.yaml && \
+                        sed -i 's/num_cpus: 0.0/num_gpus: {args.num_of_gpus}/' serve.yaml && \
                         serve run serve.yaml\"",
                 shell=True,
                 check=True,