113 kubernetes integration

README.md

@@ -68,6 +68,39 @@ segmenter.infer_single(
 )
 ```
 
+## Kubernetes Support
+BraTS orchestrator also supports Kubernetes to run the algorithms remotely, as an alternative to local execution with Docker or Singularity.
+To use Kubernetes for execution, the orchestrator will automatically use your kubeconfig file from the default location (`~/.kube/config`). If your kubeconfig file is not in the default location, set the `KUBECONFIG` environment variable to the location of your kubeconfig file:
+```bash
+export KUBECONFIG=/path/to/kubeconfig
+```
+Then, specify the backend to use as `kubernetes` (or the corresponding enum value `Backends.KUBERNETES`) when running the inference:
+```python
+from brats.constants import Backends
+segmenter.infer_single(
+    t1c="path/to/t1c.nii.gz",
+    output_file="path/to/segmentation.nii.gz",
+    backend=Backends.KUBERNETES
+)
+```
+By default, as shown above, the algorithm runs in the default Kubernetes namespace. It uses the default StorageClass and automatically creates a 1Gi PersistentVolumeClaim (PVC) to manage input and output data. If needed, you can customize settings such as the namespace, PVC name, storage size, storage class, job name, and mount path by providing related keyword arguments to the `infer_single` method. The `data_mount_path` parameter determines where the PVC will be mounted inside the Pod.
+When using Kubernetes, the algorithm is executed inside a Kubernetes Job. Input data is first uploaded to a PersistentVolume, which is mounted into the Pod running the job. After the algorithm finishes running in the Pod, the output data is transferred back from the cluster to your local machine.
+```python
+segmenter.infer_single(
+    t1c="path/to/t1c.nii.gz",
+    output_file="path/to/segmentation.nii.gz",
+    backend=Backends.KUBERNETES,
+    kubernetes_kwargs={
+        "namespace": "brats",
+        "pvc_name": "brats-iwydw55ej7qm-pvc",
+        "pvc_storage_size": "2Gi",
+        "pvc_storage_class": "brats-pvc-storage-class",
+        "job_name": "brats-oxh24nu4dhk9-job",
+        "data_mount_path": "/data",
+    }
+)
+```
+
 ## Available Algorithms and Usage
 
 > [!IMPORTANT]

brats/constants.py

@@ -13,6 +13,9 @@ class Backends(str, Enum):
     SINGULARITY = "singularity"
     """Run the algorithms using Singularity containers."""
 
+    KUBERNETES = "kubernetes"
+    """Run the algorithms using Kubernetes Jobs."""
+
 
 class Task(str, Enum):
     """Available tasks."""

brats/core/brats_algorithm.py

@@ -9,6 +9,7 @@
 
 from brats.core.docker import run_container as run_docker_container
 from brats.core.singularity import run_container as run_singularity_container
+from brats.core.kubernetes import run_job as run_kubernetes_job
 from brats.utils.algorithm_config import load_algorithms
 from brats.constants import OUTPUT_NAME_SCHEMA, Algorithms, Task, Backends
 from brats.utils.data_handling import InferenceSetup
@@ -163,6 +164,7 @@ def _get_backend_runner(self, backend: Backends) -> Optional[Callable]:
         backend_dispatch = {
             Backends.DOCKER: run_docker_container,
             Backends.SINGULARITY: run_singularity_container,
+            Backends.KUBERNETES: run_kubernetes_job,
         }
         runner = backend_dispatch.get(backend, None)
         return runner
@@ -173,6 +175,7 @@ def _infer_single(
         output_file: Path | str,
         log_file: Optional[Path | str] = None,
         backend: Backends = Backends.DOCKER,
+        kubernetes_kwargs: Optional[Dict] = None,
     ) -> None:
         """
         Perform a single inference run with the provided inputs and save the output in the specified file.
@@ -181,7 +184,8 @@ def _infer_single(
             inputs (dict[str, Path  |  str]): Input Images for the task
             output_file (Path | str): File to save the output
             log_file (Optional[Path  |  str], optional): Log file with extra information. Defaults to None.
-            backend (Backends | str, optional): Backend to use for inference. Defaults to Backends.DOCKER.
+            backend (Backends, optional): Backend to use for inference. Defaults to Backends.DOCKER.
+            kubernetes_kwargs (Optional[Dict], optional): Optional keyword arguments for Kubernetes Backend. Defaults to None.
         """
         with InferenceSetup(log_file=log_file) as (tmp_data_folder, tmp_output_folder):
             logger.info(f"Performing single inference")
@@ -199,13 +203,22 @@ def _infer_single(
             runner = self._get_backend_runner(backend)
             if runner is None:
                 raise ValueError(f"Unsupported backend: {backend}")
-            runner(
+            runner_kwargs = dict(
                 algorithm=self.algorithm,
                 data_path=tmp_data_folder,
                 output_path=tmp_output_folder,
                 cuda_devices=self.cuda_devices,
                 force_cpu=self.force_cpu,
             )
+            if kubernetes_kwargs is not None:
+                logger.debug(f"Adding Kubernetes kwargs: {kubernetes_kwargs}")
+                if backend != Backends.KUBERNETES:
+                    raise ValueError(
+                        "Kubernetes kwargs can only be used with the Kubernetes backend."
+                    )
+                for key, value in kubernetes_kwargs.items():
+                    runner_kwargs[key] = value
+            runner(**runner_kwargs)
             self._process_single_output(
                 tmp_output_folder=tmp_output_folder,
                 subject_id=subject_id,
@@ -219,6 +232,7 @@ def _infer_batch(
         output_folder: Path | str,
         log_file: Optional[Path | str] = None,
         backend: Backends = Backends.DOCKER,
+        kubernetes_kwargs: Optional[Dict] = None,
     ):
         """Perform a batch inference run with the provided inputs and save the outputs in the specified folder.
 
@@ -227,6 +241,7 @@ def _infer_batch(
             output_folder (Path | str): Folder to save the outputs
             log_file (Optional[Path  |  str], optional): Log file with extra information. Defaults to None.
             backend (Backends, optional): Backend to use for inference. Defaults to Backends.DOCKER.
+            kubernetes_kwargs (Optional[Dict], optional): Optional keyword arguments for Kubernetes Backend. Defaults to None.
         """
         with InferenceSetup(log_file=log_file) as (tmp_data_folder, tmp_output_folder):
 
@@ -246,14 +261,23 @@ def _infer_batch(
             if runner is None:
                 raise ValueError(f"Unsupported backend: {backend}")
             # run inference in container
-            runner(
+            runner_kwargs = dict(
                 algorithm=self.algorithm,
                 data_path=tmp_data_folder,
                 output_path=tmp_output_folder,
                 cuda_devices=self.cuda_devices,
                 force_cpu=self.force_cpu,
                 internal_external_name_map=internal_external_name_map,
             )
+            if kubernetes_kwargs is not None:
+                logger.debug(f"Adding Kubernetes kwargs: {kubernetes_kwargs}")
+                if backend != Backends.KUBERNETES:
+                    raise ValueError(
+                        "Kubernetes kwargs can only be used with the Kubernetes backend."
+                    )
+                for key, value in kubernetes_kwargs.items():
+                    runner_kwargs[key] = value
+            runner(**runner_kwargs)
 
             self._process_batch_output(
                 tmp_output_folder=tmp_output_folder,