feat(cli): allow multi-gpus when running inference locally

instill-ai · Dec 6, 2024 · f4d25cf · f4d25cf
1 parent e38eefe
commit f4d25cf
Showing 1 changed file with 9 additions and 1 deletion.
diff --git 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,