Add --gpu to set machine_shape for push

Author: stevemessick

src/kaggle/api/kaggle_api_extended.py

@@ -3038,7 +3038,9 @@ def kernels_initialize_cli(self, folder=None):
         meta_file = self.kernels_initialize(folder)
         print("Kernel metadata template written to: " + meta_file)
 
-    def kernels_push(self, folder: str, timeout: Optional[str] = None) -> ApiSaveKernelResponse:
+    def kernels_push(
+        self, folder: str, timeout: Optional[str] = None, gpu: Optional[str] = None
+    ) -> ApiSaveKernelResponse:
         """Pushes a kernel to Kaggle.
 
         This method reads the metadata file and kernel files from a notebook,
@@ -3047,6 +3049,8 @@ def kernels_push(self, folder: str, timeout: Optional[str] = None) -> ApiSaveKer
         Args:
             folder (str): The path to the folder.
             timeout (Optional[str]): The maximum run time in seconds.
+            gpu (Optional[str]): The type of GPU to use for the kernel run. If set, this value overrides boolean
+                settings for GPU/TPU found int metadata.
 
         Returns:
             ApiSaveKernelResponse: An ApiSaveKernelResponse object.
@@ -3165,6 +3169,9 @@ def kernels_push(self, folder: str, timeout: Optional[str] = None) -> ApiSaveKer
             request.docker_image = self.get_or_default(meta_data, "docker_image", None)
             if timeout:
                 request.session_timeout_seconds = int(timeout)
+            if gpu:
+                # The allowed names are in an enum that is not currently included in kagglesdk.
+                request.machine_shape = gpu
             # Without the type hint, mypy thinks save_kernel() has type Any when checking warn_return_any.
             response: ApiSaveKernelResponse = kaggle.kernels.kernels_api_client.save_kernel(request)
             return response

src/kaggle/cli.py

@@ -550,6 +550,7 @@ def parse_kernels(subparsers) -> None:
     parser_kernels_push_optional.add_argument(
         "-t", "--timeout", type=int, dest="timeout", help=Help.param_kernel_timeout
     )
+    parser_kernels_push_optional.add_argument("--gpu", dest="gpu", help=Help.param_kernel_gpu)
     parser_kernels_push._action_groups.append(parser_kernels_push_optional)
     parser_kernels_push.set_defaults(func=api.kernels_push_cli)
 
@@ -1299,6 +1300,7 @@ class Help(object):
     param_kernel_output_file_pattern = (
         "Regex pattern to match against filenames. Only files matching the pattern will be downloaded."
     )
+    param_kernel_gpu = "Specify the type of GPU to use for the kernel run"
 
     # Models params
     param_model = "Model URL suffix in format <owner>/<model-name>"