Add GPU support for inference service

app/services/inference.py

@@ -23,7 +23,8 @@ class ResNetInferenceService(InferenceService[ImageRequest, PredictionResponse])
         self.processor = None
         self._is_loaded = False
         self.model_path = os.path.join("models", model_name)
-        logger.info(f"Initializing ResNet service: {self.model_path}")
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        logger.info(f"Initializing ResNet service: {self.model_path} on {self.device}")
 
     def load_model(self) -> None:
         if self._is_loaded:
@@ -48,6 +49,7 @@ class ResNetInferenceService(InferenceService[ImageRequest, PredictionResponse])
                 self.model_path, local_files_only=True
             )
             assert self.model is not None
+            self.model.to(self.device)
 
         self._is_loaded = True
         logger.info(f"Model loaded: {len(self.model.config.id2label)} classes")  # pyright: ignore
@@ -65,6 +67,7 @@ class ResNetInferenceService(InferenceService[ImageRequest, PredictionResponse])
             image = image.convert('RGB')
 
         inputs = self.processor(image, return_tensors="pt")
+        inputs = inputs.to(self.device)
 
         with torch.no_grad():
             logits = self.model(**inputs).logits.squeeze()   # pyright: ignore

challenge-cli.py

@@ -114,6 +114,12 @@ def cli():
     " You can pass the model.id from a previous invocation.",
     metavar="ID",
 )
+@click.option(
+    "--gpu",
+    is_flag=True,
+    default=False,
+    help="Request a GPU (NVIDIA L4) for the inference service.",
+)
 def upload_submission(
     account: str,
     name: str,
@@ -123,6 +129,7 @@ def upload_submission(
     volume_mount: Path | None,
     artifact_id: str | None,
     model_id: str | None,
+    gpu: bool,
 ) -> None:
     dyffapi = Client()
 
@@ -208,6 +215,16 @@ def upload_submission(
     else:
         volumeMounts = None
 
+    accelerator: Accelerator | None = None
+    if gpu:
+        accelerator = Accelerator(
+            kind="GPU",
+            gpu=AcceleratorGPU(
+                hardwareTypes=["nvidia-l4"],
+                count=1,
+            ),
+        )
+
     # Don't change this
     service_request = InferenceServiceCreateRequest(
         account=account,
@@ -218,6 +235,7 @@ def upload_submission(
             imageRef=EntityIdentifier.of(artifact),
             resources=ModelResources(),
             volumeMounts=volumeMounts,
+            accelerator=accelerator,
         ),
         interface=InferenceInterface(
             endpoint=endpoint,