feat: Add cli for benchmark tasks

Author: dolaameng

pyproject.toml

@@ -32,6 +32,7 @@ dependencies = [
     "urllib3 >= 1.15.1",
     "packaging",
     "protobuf",
+    "jupytext >= 1.15.0",
 ]
 
 [project.scripts]

src/kaggle/api/kaggle_api.py

@@ -338,7 +338,7 @@ def kernel_push(self, kernel_push_request):  # noqa: E501
         """
         with tempfile.TemporaryDirectory() as tmpdir:
             meta_file = os.path.join(tmpdir, "kernel-metadata.json")
-            (fd, code_file) = tempfile.mkstemp("code", "py", tmpdir, text=True)
+            fd, code_file = tempfile.mkstemp("code", "py", tmpdir, text=True)
             fd.write(json.dumps(kernel_push_request.code))
             os.close(fd)
             with open(meta_file, "w") as f:

src/kaggle/api/kaggle_api_extended.py

@@ -721,7 +721,7 @@ def _authenticate_with_legacy_apikey(self) -> bool:
         return True
 
     def _authenticate_with_access_token(self):
-        (access_token, source) = get_access_token_from_env()
+        access_token, source = get_access_token_from_env()
         if not access_token:
             return False
 
@@ -1901,7 +1901,7 @@ def dataset_metadata_update(self, dataset, path):
             dataset: The dataset to update.
             path: The path to the metadata file.
         """
-        (owner_slug, dataset_slug, effective_path) = self.dataset_metadata_prep(dataset, path)
+        owner_slug, dataset_slug, effective_path = self.dataset_metadata_prep(dataset, path)
         meta_file = self.get_dataset_metadata_file(effective_path)
         with open(meta_file, "r") as f:
             metadata = json.load(f)
@@ -1954,7 +1954,7 @@ def dataset_metadata(self, dataset, path):
         Returns:
             The path to the downloaded metadata file.
         """
-        (owner_slug, dataset_slug, effective_path) = self.dataset_metadata_prep(dataset, path)
+        owner_slug, dataset_slug, effective_path = self.dataset_metadata_prep(dataset, path)
 
         if not os.path.exists(effective_path):
             os.makedirs(effective_path)
@@ -3462,7 +3462,7 @@ def kernels_output_cli(self, kernel, kernel_opt=None, path=None, force=False, qu
             file_pattern: Regex pattern to match against filenames. Only files matching the pattern will be downloaded.
         """
         kernel = kernel or kernel_opt
-        (_, token) = self.kernels_output(kernel, path, file_pattern, force, quiet)
+        _, token = self.kernels_output(kernel, path, file_pattern, force, quiet)
         if token:
             print(f"Next page token: {token}")
 
@@ -3508,6 +3508,217 @@ def kernels_status_cli(self, kernel, kernel_opt=None):
         else:
             print('%s has status "%s"' % (kernel, status))
 
+    def benchmarks_pull(self, kernel: str, path: str = None, quiet: bool = False):
+        """Pulls a benchmark notebook and converts it to a .py script."""
+        import os
+
+        try:
+            import jupytext
+        except ImportError:
+            raise ImportError("jupytext is required for benchmarks functionality. Please install it.")
+
+        effective_path = self.kernels_pull(kernel, path=path, metadata=True, quiet=quiet)
+
+        # After pulling, find the .ipynb file in effective_path
+        ipynb_files = [f for f in os.listdir(effective_path) if f.endswith(".ipynb")]
+        if not ipynb_files:
+            raise ValueError("Could not find a .ipynb file in the pulled kernel.")
+
+        # Rename the first .ipynb file to benchmark.ipynb (if it isn't already)
+        ipynb_name = ipynb_files[0]
+        ipynb_path = os.path.join(effective_path, ipynb_name)
+        benchmark_ipynb_path = os.path.join(effective_path, "benchmark.ipynb")
+        if ipynb_name != "benchmark.ipynb":
+            os.rename(ipynb_path, benchmark_ipynb_path)
+            if not quiet:
+                print(f"Renamed {ipynb_name} to benchmark.ipynb")
+
+        # Convert to benchmark.py
+        benchmark_py_path = os.path.join(effective_path, "benchmark.py")
+        notebook = jupytext.read(benchmark_ipynb_path)
+        
+        # Strip confusing metadata formatting from the resulting Python document
+        if "jupytext" not in notebook.metadata:
+            notebook.metadata["jupytext"] = {}
+        notebook.metadata["jupytext"]["notebook_metadata_filter"] = "-all"
+        notebook.metadata["jupytext"]["cell_metadata_filter"] = "-all"
+        
+        jupytext.write(notebook, benchmark_py_path, fmt="py:percent")
+        if not quiet:
+            print(f"Converted benchmark.ipynb to {benchmark_py_path}")
+        return effective_path
+
+    def benchmarks_pull_cli(self, kernel, kernel_opt=None, path=None):
+        kernel = kernel or kernel_opt
+        effective_path = self.benchmarks_pull(kernel, path=path, quiet=False)
+        print(f"Benchmark pulled and converted successfully to {effective_path}")
+
+    def benchmarks_publish_and_run(
+        self, kernel: str = None, path: str = None, file_name: str = None, quiet: bool = False
+    ):
+        """Converts a local .py benchmark to .ipynb and pushes it to Kaggle."""
+        import os
+        import json
+
+        try:
+            import jupytext
+        except ImportError:
+            raise ImportError("jupytext is required for benchmarks functionality. Please install it.")
+
+        path = path or os.getcwd()
+        file_name = file_name or "benchmark.py"
+        py_path = os.path.join(path, file_name)
+
+        if not os.path.exists(py_path):
+            raise FileNotFoundError(f"Source file not found: {py_path}")
+
+        ipynb_path = os.path.join(path, "benchmark.ipynb")
+        notebook = jupytext.read(py_path, fmt="py:percent")
+        
+        # Inject default Kaggle kernelspec so Papermill executes the notebook correctly
+        if "kernelspec" not in notebook.metadata:
+            notebook.metadata["kernelspec"] = {
+                "display_name": "Python 3",
+                "language": "python",
+                "name": "python3"
+            }
+
+        jupytext.write(notebook, ipynb_path)
+        if not quiet:
+            print(f"Converted {py_path} to {ipynb_path}")
+
+        # Ensure kernel-metadata.json exists and has "personal-benchmark"
+        metadata_path = os.path.join(path, self.KERNEL_METADATA_FILE)
+        if not os.path.exists(metadata_path):
+            # Create a default metadata file
+            if not kernel:
+                raise ValueError("A kernel slug must be specified to create a new metadata file.")
+
+            if "/" in kernel:
+                self.validate_kernel_string(kernel)
+                owner_slug, kernel_slug = kernel.split("/")
+            else:
+                owner_slug = self.get_config_value(self.CONFIG_NAME_USER)
+                kernel_slug = kernel
+
+            title = kernel_slug.replace("-", " ").title()
+            metadata = {
+                "id": f"{owner_slug}/{kernel_slug}",
+                "title": title,
+                "code_file": "benchmark.ipynb",
+                "language": "python",
+                "kernel_type": "notebook",
+                "is_private": "true",
+                "enable_gpu": "false",
+                "enable_internet": "true",
+                "dataset_sources": [],
+                "competition_sources": [],
+                "kernel_sources": [],
+                "model_sources": [],
+                "keywords": ["personal-benchmark"],
+            }
+            with open(metadata_path, "w") as f:
+                json.dump(metadata, f, indent=2)
+            if not quiet:
+                print(f"Created kernel metadata at {metadata_path}")
+        else:
+            # Read existing and inject keyword if missing
+            with open(metadata_path, "r") as f:
+                metadata = json.load(f)
+
+            if kernel:
+                if "/" in kernel:
+                    self.validate_kernel_string(kernel)
+                    owner_slug, kernel_slug = kernel.split("/")
+                else:
+                    owner_slug = self.get_config_value(self.CONFIG_NAME_USER)
+                    kernel_slug = kernel
+
+                new_id = f"{owner_slug}/{kernel_slug}"
+                if metadata.get("id") != new_id:
+                    metadata["id"] = new_id
+                    metadata["title"] = kernel_slug.replace("-", " ").title()
+                    if "id_no" in metadata:
+                        del metadata["id_no"]
+
+            if "keywords" not in metadata:
+                metadata["keywords"] = []
+            if "personal-benchmark" not in metadata["keywords"]:
+                metadata["keywords"].append("personal-benchmark")
+            if "code_file" not in metadata or metadata["code_file"] != "benchmark.ipynb":
+                metadata["code_file"] = "benchmark.ipynb"
+
+            with open(metadata_path, "w") as f:
+                json.dump(metadata, f, indent=2)
+
+        # Now push using kernels_push
+        return self.kernels_push(path)
+
+    def benchmarks_publish_and_run_cli(self, kernel=None, kernel_opt=None, path=None, file_name=None):
+        kernel = kernel or kernel_opt
+        result = self.benchmarks_publish_and_run(kernel, path=path, file_name=file_name, quiet=False)
+        
+        url_text = ""
+        if result and getattr(result, "url", None):
+            url_text = f"\nTracking URL: {result.url}"
+            
+        print(f"Benchmark pushed and started successfully.{url_text}\nRun kaggle benchmarks tasks results to stream output.")
+
+    def benchmarks_get_results(self, kernel: str = None, path: str = None, poll_interval: int = 60, timeout: int = None):
+        """Polls the status of a benchmark until complete, then downloads the output."""
+        import os
+        import time
+        import json
+
+        if kernel is None:
+            check_path = path or os.getcwd()
+            meta_file = os.path.join(check_path, self.KERNEL_METADATA_FILE)
+            if os.path.exists(meta_file):
+                with open(meta_file, "r") as f:
+                    try:
+                        metadata = json.load(f)
+                        kernel = metadata.get("id")
+                    except Exception:
+                        pass
+        if kernel is None:
+            raise ValueError("A kernel must be specified")
+
+        start_time = time.time()
+        print(f"Waiting for benchmark {kernel} to complete...")
+
+        while True:
+            response = self.kernels_status(kernel)
+            status = response.status
+            status_str = str(status).upper()
+            if "COMPLETE" in status_str:
+                print(f"Benchmark {kernel} completed.")
+                break
+            elif "ERROR" in status_str:
+                message = getattr(response, "failure_message", "")
+                error_txt = f" Message: {message}" if message else ""
+                print(f"Benchmark {kernel} failed!{error_txt}")
+                print(f"Attempting to download partial logs for debugging...")
+                try:
+                    self.kernels_output(kernel, path=check_path, force=True, quiet=False)
+                except Exception as log_err:
+                    print(f"Could not retrieve backend logs: {log_err}")
+                raise ValueError(f"Benchmark execution terminated with an error state.")
+            else:
+                print(f"Status: {status}. Waiting {poll_interval}s...")
+                time.sleep(poll_interval)
+
+            if timeout is not None and (time.time() - start_time) > timeout:
+                raise TimeoutError(f"Timed out waiting for benchmark after {timeout} seconds.")
+
+        # Now download output
+        print(f"Downloading results for {kernel}...")
+        return self.kernels_output(kernel=kernel, path=path, force=True, quiet=False)
+
+    def benchmarks_get_results_cli(self, kernel, kernel_opt=None, path=None, poll_interval=60, timeout=None):
+        kernel = kernel or kernel_opt
+        self.benchmarks_get_results(kernel, path=path, poll_interval=poll_interval, timeout=timeout)
+        print("Output downloaded successfully.")
+
     def model_get(self, model: str) -> ApiModel:
         """Gets a model.
 
@@ -4587,7 +4798,7 @@ def files_upload_cli(self, local_paths, inbox_path, no_resume, no_compress):
         files_to_create = []
         with ResumableUploadContext(no_resume) as upload_context:
             for local_path in local_paths:
-                (upload_file, file_name) = self.file_upload_cli(local_path, inbox_path, no_compress, upload_context)
+                upload_file, file_name = self.file_upload_cli(local_path, inbox_path, no_compress, upload_context)
                 if upload_file is None:
                     continue