mlrun 1.7.2rc3__py3-none-any.whl → 1.8.0rc2__py3-none-any.whl

mlrun/runtimes/nuclio/application/application.py

@@ -275,7 +275,7 @@ class ApplicationRuntime(RemoteRuntime):
         tag="",
         verbose=False,
         auth_info: schemas.AuthInfo = None,
-        builder_env: dict = None,
+        builder_env: typing.Optional[dict] = None,
         force_build: bool = False,
         with_mlrun=None,
         skip_deployed=False,
@@ -367,7 +367,7 @@ class ApplicationRuntime(RemoteRuntime):
         source,
         workdir=None,
         pull_at_runtime: bool = False,
-        target_dir: str = None,
+        target_dir: typing.Optional[str] = None,
     ):
         """load the code from git/tar/zip archive at build
 
@@ -426,12 +426,12 @@ class ApplicationRuntime(RemoteRuntime):
 
     def create_api_gateway(
         self,
-        name: str = None,
-        path: str = None,
+        name: typing.Optional[str] = None,
+        path: typing.Optional[str] = None,
         direct_port_access: bool = False,
         authentication_mode: schemas.APIGatewayAuthenticationMode = None,
-        authentication_creds: tuple[str, str] = None,
-        ssl_redirect: bool = None,
+        authentication_creds: typing.Optional[tuple[str, str]] = None,
+        ssl_redirect: typing.Optional[bool] = None,
         set_as_default: bool = False,
         gateway_timeout: typing.Optional[int] = None,
     ):
@@ -540,13 +540,13 @@ class ApplicationRuntime(RemoteRuntime):
         self,
         path: str = "",
         body: typing.Optional[typing.Union[str, bytes, dict]] = None,
-        method: str = None,
-        headers: dict = None,
+        method: typing.Optional[str] = None,
+        headers: typing.Optional[dict] = None,
         dashboard: str = "",
         force_external_address: bool = False,
         auth_info: schemas.AuthInfo = None,
-        mock: bool = None,
-        credentials: tuple[str, str] = None,
+        mock: typing.Optional[bool] = None,
+        credentials: typing.Optional[tuple[str, str]] = None,
         **http_client_kwargs,
     ):
         self._sync_api_gateway()
@@ -653,7 +653,7 @@ class ApplicationRuntime(RemoteRuntime):
 
     def _build_application_image(
         self,
-        builder_env: dict = None,
+        builder_env: typing.Optional[dict] = None,
         force_build: bool = False,
         watch=True,
         with_mlrun=None,

mlrun/runtimes/nuclio/function.py

@@ -26,9 +26,6 @@ import requests
 import semver
 from aiohttp.client import ClientSession
 from kubernetes import client
-from mlrun_pipelines.common.mounts import VolumeMount
-from mlrun_pipelines.common.ops import deploy_op
-from mlrun_pipelines.mounts import mount_v3io, v3io_cred
 from nuclio.deploy import find_dashboard_url, get_deploy_status
 from nuclio.triggers import V3IOStreamTrigger
 
@@ -47,9 +44,11 @@ from mlrun.platforms.iguazio import (
     split_path,
 )
 from mlrun.runtimes.base import FunctionStatus, RunError
+from mlrun.runtimes.mounts import VolumeMount, mount_v3io, v3io_cred
 from mlrun.runtimes.pod import KubeResource, KubeResourceSpec
 from mlrun.runtimes.utils import get_item_name, log_std
 from mlrun.utils import get_in, logger, update_in
+from mlrun_pipelines.common.ops import deploy_op
 
 
 def validate_nuclio_version_compatibility(*min_versions):
@@ -554,7 +553,7 @@ class RemoteRuntime(KubeResource):
         tag="",
         verbose=False,
         auth_info: AuthInfo = None,
-        builder_env: dict = None,
+        builder_env: typing.Optional[dict] = None,
         force_build: bool = False,
     ):
         """Deploy the nuclio function to the cluster
@@ -693,7 +692,9 @@ class RemoteRuntime(KubeResource):
         super().with_priority_class(name)
 
     def with_service_type(
-        self, service_type: str, add_templated_ingress_host_mode: str = None
+        self,
+        service_type: str,
+        add_templated_ingress_host_mode: typing.Optional[str] = None,
     ):
         """
         Enables to control the service type of the pod and the addition of templated ingress host
@@ -887,13 +888,13 @@ class RemoteRuntime(KubeResource):
     def invoke(
         self,
         path: str,
-        body: typing.Union[str, bytes, dict] = None,
-        method: str = None,
-        headers: dict = None,
+        body: typing.Optional[typing.Union[str, bytes, dict]] = None,
+        method: typing.Optional[str] = None,
+        headers: typing.Optional[dict] = None,
         dashboard: str = "",
         force_external_address: bool = False,
         auth_info: AuthInfo = None,
-        mock: bool = None,
+        mock: typing.Optional[bool] = None,
         **http_client_kwargs,
     ):
         """Invoke the remote (live) function and return the results
@@ -995,8 +996,8 @@ class RemoteRuntime(KubeResource):
 
     def with_sidecar(
         self,
-        name: str = None,
-        image: str = None,
+        name: typing.Optional[str] = None,
+        image: typing.Optional[str] = None,
         ports: typing.Optional[typing.Union[int, list[int]]] = None,
         command: typing.Optional[str] = None,
         args: typing.Optional[list[str]] = None,
@@ -1035,10 +1036,9 @@ class RemoteRuntime(KubeResource):
         if args and sidecar.get("command"):
             sidecar["args"] = mlrun.utils.helpers.as_list(args)
 
-        # put the configured resources on the sidecar container instead of the reverse proxy container
+        # populate the sidecar resources from the function spec
         if self.spec.resources:
             sidecar["resources"] = self.spec.resources
-            self.spec.resources = None
 
     def _set_sidecar(self, name: str) -> dict:
         self.spec.config.setdefault("spec.sidecars", [])

mlrun/runtimes/nuclio/serving.py

@@ -51,7 +51,7 @@ if TYPE_CHECKING:
 def new_v2_model_server(
     name,
     model_class: str,
-    models: dict = None,
+    models: Optional[dict] = None,
     filename="",
     protocol="",
     image="",
@@ -356,12 +356,12 @@ class ServingRuntime(RemoteRuntime):
     def add_model(
         self,
         key: str,
-        model_path: str = None,
-        class_name: str = None,
-        model_url: str = None,
-        handler: str = None,
-        router_step: str = None,
-        child_function: str = None,
+        model_path: Optional[str] = None,
+        class_name: Optional[str] = None,
+        model_url: Optional[str] = None,
+        handler: Optional[str] = None,
+        router_step: Optional[str] = None,
+        child_function: Optional[str] = None,
         **class_args,
     ):
         """add ml model and/or route to the function.
@@ -509,7 +509,7 @@ class ServingRuntime(RemoteRuntime):
                         stream.path, group=group, shards=stream.shards, **trigger_args
                     )
 
-    def _deploy_function_refs(self, builder_env: dict = None):
+    def _deploy_function_refs(self, builder_env: Optional[dict] = None):
         """set metadata and deploy child functions"""
         for function_ref in self._spec.function_refs.values():
             logger.info(f"deploy child function {function_ref.name} ...")
@@ -583,7 +583,7 @@ class ServingRuntime(RemoteRuntime):
         tag="",
         verbose=False,
         auth_info: mlrun.common.schemas.AuthInfo = None,
-        builder_env: dict = None,
+        builder_env: Optional[dict] = None,
         force_build: bool = False,
     ):
         """deploy model serving function to a local/remote cluster

mlrun/runtimes/pod.py

@@ -17,14 +17,16 @@ import os
 import re
 import time
 import typing
+from collections.abc import Iterable
 from enum import Enum
 
 import dotenv
 import kubernetes.client as k8s_client
-import mlrun_pipelines.mounts
-from mlrun_pipelines.mixins import KfpAdapterMixin
+from kubernetes.client import V1Volume, V1VolumeMount
 
+import mlrun.common.constants
 import mlrun.errors
+import mlrun.runtimes.mounts
 import mlrun.utils.regex
 from mlrun.common.schemas import (
     NodeSelectorOperator,
@@ -309,7 +311,7 @@ class KubeResourceSpec(FunctionSpec):
         return self._termination_grace_period_seconds
 
     def _serialize_field(
-        self, struct: dict, field_name: str = None, strip: bool = False
+        self, struct: dict, field_name: typing.Optional[str] = None, strip: bool = False
     ) -> typing.Any:
         """
         Serialize a field to a dict, list, or primitive type.
@@ -321,7 +323,7 @@ class KubeResourceSpec(FunctionSpec):
         return super()._serialize_field(struct, field_name, strip)
 
     def _enrich_field(
-        self, struct: dict, field_name: str = None, strip: bool = False
+        self, struct: dict, field_name: typing.Optional[str] = None, strip: bool = False
     ) -> typing.Any:
         k8s_api = k8s_client.ApiClient()
         if strip:
@@ -366,6 +368,35 @@ class KubeResourceSpec(FunctionSpec):
                 + f"service accounts {allowed_service_accounts}"
             )
 
+    def with_volumes(
+        self,
+        volumes: typing.Union[list[dict], dict, V1Volume],
+    ) -> "KubeResourceSpec":
+        """Add volumes to the volumes dictionary, only used as part of the mlrun_pipelines mount functions."""
+        if isinstance(volumes, dict):
+            set_named_item(self._volumes, volumes)
+        elif isinstance(volumes, Iterable):
+            for volume in volumes:
+                set_named_item(self._volumes, volume)
+        else:
+            set_named_item(self._volumes, volumes)
+        return self
+
+    def with_volume_mounts(
+        self,
+        volume_mounts: typing.Union[list[dict], dict, V1VolumeMount],
+    ) -> "KubeResourceSpec":
+        """Add volume mounts to the volume mounts dictionary,
+        only used as part of the mlrun_pipelines mount functions."""
+        if isinstance(volume_mounts, dict):
+            self._set_volume_mount(volume_mounts)
+        elif isinstance(volume_mounts, Iterable):
+            for volume_mount in volume_mounts:
+                self._set_volume_mount(volume_mount)
+        else:
+            self._set_volume_mount(volume_mounts)
+        return self
+
     def _set_volume_mount(
         self, volume_mount, volume_mounts_field_name="_volume_mounts"
     ):
@@ -380,9 +411,9 @@ class KubeResourceSpec(FunctionSpec):
     def _verify_and_set_limits(
         self,
         resources_field_name,
-        mem: str = None,
-        cpu: str = None,
-        gpus: int = None,
+        mem: typing.Optional[str] = None,
+        cpu: typing.Optional[str] = None,
+        gpus: typing.Optional[int] = None,
         gpu_type: str = "nvidia.com/gpu",
         patch: bool = False,
     ):
@@ -430,8 +461,8 @@ class KubeResourceSpec(FunctionSpec):
     def _verify_and_set_requests(
         self,
         resources_field_name,
-        mem: str = None,
-        cpu: str = None,
+        mem: typing.Optional[str] = None,
+        cpu: typing.Optional[str] = None,
         patch: bool = False,
     ):
         resources = verify_requests(resources_field_name, mem=mem, cpu=cpu)
@@ -456,9 +487,9 @@ class KubeResourceSpec(FunctionSpec):
 
     def with_limits(
         self,
-        mem: str = None,
-        cpu: str = None,
-        gpus: int = None,
+        mem: typing.Optional[str] = None,
+        cpu: typing.Optional[str] = None,
+        gpus: typing.Optional[int] = None,
         gpu_type: str = "nvidia.com/gpu",
         patch: bool = False,
     ):
@@ -474,7 +505,12 @@ class KubeResourceSpec(FunctionSpec):
         """
         self._verify_and_set_limits("resources", mem, cpu, gpus, gpu_type, patch=patch)
 
-    def with_requests(self, mem: str = None, cpu: str = None, patch: bool = False):
+    def with_requests(
+        self,
+        mem: typing.Optional[str] = None,
+        cpu: typing.Optional[str] = None,
+        patch: bool = False,
+    ):
         """
         Set requested (desired) pod cpu/memory resources
 
@@ -936,12 +972,12 @@ class AutoMountType(str, Enum):
     @classmethod
     def all_mount_modifiers(cls):
         return [
-            mlrun_pipelines.mounts.v3io_cred.__name__,
-            mlrun_pipelines.mounts.mount_v3io.__name__,
-            mlrun_pipelines.mounts.mount_pvc.__name__,
-            mlrun_pipelines.mounts.auto_mount.__name__,
-            mlrun_pipelines.mounts.mount_s3.__name__,
-            mlrun_pipelines.mounts.set_env_variables.__name__,
+            mlrun.runtimes.mounts.v3io_cred.__name__,
+            mlrun.runtimes.mounts.mount_v3io.__name__,
+            mlrun.runtimes.mounts.mount_pvc.__name__,
+            mlrun.runtimes.mounts.auto_mount.__name__,
+            mlrun.runtimes.mounts.mount_s3.__name__,
+            mlrun.runtimes.mounts.set_env_variables.__name__,
         ]
 
     @classmethod
@@ -958,27 +994,27 @@ class AutoMountType(str, Enum):
     def _get_auto_modifier():
         # If we're running on Iguazio - use v3io_cred
         if mlconf.igz_version != "":
-            return mlrun_pipelines.mounts.v3io_cred
+            return mlrun.runtimes.mounts.v3io_cred
         # Else, either pvc mount if it's configured or do nothing otherwise
         pvc_configured = (
             "MLRUN_PVC_MOUNT" in os.environ
             or "pvc_name" in mlconf.get_storage_auto_mount_params()
         )
-        return mlrun_pipelines.mounts.mount_pvc if pvc_configured else None
+        return mlrun.runtimes.mounts.mount_pvc if pvc_configured else None
 
     def get_modifier(self):
         return {
             AutoMountType.none: None,
-            AutoMountType.v3io_credentials: mlrun_pipelines.mounts.v3io_cred,
-            AutoMountType.v3io_fuse: mlrun_pipelines.mounts.mount_v3io,
-            AutoMountType.pvc: mlrun_pipelines.mounts.mount_pvc,
+            AutoMountType.v3io_credentials: mlrun.runtimes.mounts.v3io_cred,
+            AutoMountType.v3io_fuse: mlrun.runtimes.mounts.mount_v3io,
+            AutoMountType.pvc: mlrun.runtimes.mounts.mount_pvc,
             AutoMountType.auto: self._get_auto_modifier(),
-            AutoMountType.s3: mlrun_pipelines.mounts.mount_s3,
-            AutoMountType.env: mlrun_pipelines.mounts.set_env_variables,
+            AutoMountType.s3: mlrun.runtimes.mounts.mount_s3,
+            AutoMountType.env: mlrun.runtimes.mounts.set_env_variables,
         }[self]
 
 
-class KubeResource(BaseRuntime, KfpAdapterMixin):
+class KubeResource(BaseRuntime):
     """
     A parent class for runtimes that generate k8s resources when executing.
     """
@@ -1020,7 +1056,8 @@ class KubeResource(BaseRuntime, KfpAdapterMixin):
 
     def get_env(self, name, default=None):
         """Get the pod environment variable for the given name, if not found return the default
-        If it's a scalar value, will return it, if the value is from source, return the k8s struct (V1EnvVarSource)"""
+        If it's a scalar value, will return it, if the value is from source, return the k8s struct (V1EnvVarSource)
+        """
         for env_var in self.spec.env:
             if get_item_name(env_var) == name:
                 # valueFrom is a workaround for now, for some reason the envs aren't getting sanitized
@@ -1050,7 +1087,11 @@ class KubeResource(BaseRuntime, KfpAdapterMixin):
         self.spec.env.append(new_var)
         return self
 
-    def set_envs(self, env_vars: dict = None, file_path: str = None):
+    def set_envs(
+        self,
+        env_vars: typing.Optional[dict] = None,
+        file_path: typing.Optional[str] = None,
+    ):
         """set pod environment var from key/value dict or .env file
 
         :param env_vars:  dict with env key/values
@@ -1070,11 +1111,16 @@ class KubeResource(BaseRuntime, KfpAdapterMixin):
             else:
                 raise mlrun.errors.MLRunNotFoundError(f"{file_path} does not exist")
         for name, value in env_vars.items():
-            self.set_env(name, value)
+            if isinstance(value, dict) and "valueFrom" in value:
+                self.set_env(name, value_from=value["valueFrom"])
+            else:
+                self.set_env(name, value)
         return self
 
     def set_image_pull_configuration(
-        self, image_pull_policy: str = None, image_pull_secret_name: str = None
+        self,
+        image_pull_policy: typing.Optional[str] = None,
+        image_pull_secret_name: typing.Optional[str] = None,
     ):
         """
         Configure the image pull parameters for the runtime.
@@ -1123,9 +1169,9 @@ class KubeResource(BaseRuntime, KfpAdapterMixin):
 
     def with_limits(
         self,
-        mem: str = None,
-        cpu: str = None,
-        gpus: int = None,
+        mem: typing.Optional[str] = None,
+        cpu: typing.Optional[str] = None,
+        gpus: typing.Optional[int] = None,
         gpu_type: str = "nvidia.com/gpu",
         patch: bool = False,
     ):
@@ -1141,7 +1187,12 @@ class KubeResource(BaseRuntime, KfpAdapterMixin):
         """
         self.spec.with_limits(mem, cpu, gpus, gpu_type, patch=patch)
 
-    def with_requests(self, mem: str = None, cpu: str = None, patch: bool = False):
+    def with_requests(
+        self,
+        mem: typing.Optional[str] = None,
+        cpu: typing.Optional[str] = None,
+        patch: bool = False,
+    ):
         """
         Set requested (desired) pod cpu/memory resources
 
@@ -1253,6 +1304,36 @@ class KubeResource(BaseRuntime, KfpAdapterMixin):
             )
         self.spec.security_context = security_context
 
+    def apply(
+        self,
+        modifier: typing.Callable[["KubeResource"], "KubeResource"],
+    ) -> "KubeResource":
+        """
+        Apply a modifier to the runtime which is used to change the runtimes k8s object's spec.
+        All modifiers accept Kube, apply some changes on its spec and return it so modifiers can be chained
+        one after the other.
+
+        :param modifier: a modifier callable object
+        :return: the runtime (self) after the modifications
+        """
+        modifier(self)
+        if AutoMountType.is_auto_modifier(modifier):
+            self.spec.disable_auto_mount = True
+
+        api_client = k8s_client.ApiClient()
+        if self.spec.env:
+            for index, env in enumerate(
+                api_client.sanitize_for_serialization(self.spec.env)
+            ):
+                self.spec.env[index] = env
+
+        if self.spec.volumes and self.spec.volume_mounts:
+            vols = api_client.sanitize_for_serialization(self.spec.volumes)
+            mounts = api_client.sanitize_for_serialization(self.spec.volume_mounts)
+            self.spec.update_vols_and_mounts(vols, mounts)
+
+        return self
+
     def list_valid_priority_class_names(self):
         return mlconf.get_valid_function_priority_class_names()
 
@@ -1412,20 +1493,32 @@ class KubeResource(BaseRuntime, KfpAdapterMixin):
     ):
         db = self._get_db()
         offset = 0
+        events_offset = 0
         try:
-            text, _ = db.get_builder_status(self, 0, logs=logs)
+            text, _, deploy_status_text_kind = db.get_builder_status(
+                self,
+                offset=0,
+                logs=logs,
+                events_offset=0,
+            )
         except mlrun.db.RunDBError:
-            raise ValueError("function or build process not found")
+            raise ValueError("Function or build process not found")
 
-        def print_log(text):
-            if text and (
+        def print_log(_text):
+            if _text and (
                 not show_on_failure
                 or self.status.state == mlrun.common.schemas.FunctionState.error
             ):
-                print(text, end="")
+                print(_text, end="")
 
         print_log(text)
-        offset += len(text)
+        if (
+            deploy_status_text_kind
+            == mlrun.common.constants.DeployStatusTextKind.events
+        ):
+            events_offset += len(text)
+        else:
+            offset += len(text)
         if watch:
             while self.status.state in [
                 mlrun.common.schemas.FunctionState.pending,
@@ -1434,14 +1527,30 @@ class KubeResource(BaseRuntime, KfpAdapterMixin):
                 time.sleep(2)
                 if show_on_failure:
                     text = ""
-                    db.get_builder_status(self, 0, logs=False)
+                    db.get_builder_status(self, offset=0, logs=False, events_offset=0)
                     if self.status.state == mlrun.common.schemas.FunctionState.error:
                         # re-read the full log on failure
-                        text, _ = db.get_builder_status(self, offset, logs=logs)
+                        text, _, deploy_status_text_kind = db.get_builder_status(
+                            self,
+                            offset=offset,
+                            logs=logs,
+                            events_offset=events_offset,
+                        )
                 else:
-                    text, _ = db.get_builder_status(self, offset, logs=logs)
+                    text, _, deploy_status_text_kind = db.get_builder_status(
+                        self,
+                        offset=offset,
+                        logs=logs,
+                        events_offset=events_offset,
+                    )
                 print_log(text)
-                offset += len(text)
+                if (
+                    deploy_status_text_kind
+                    == mlrun.common.constants.DeployStatusTextKind.events
+                ):
+                    events_offset += len(text)
+                else:
+                    offset += len(text)
 
         return self.status.state
 

mlrun/runtimes/remotesparkjob.py

@@ -13,12 +13,13 @@
 # limitations under the License.
 import re
 from subprocess import run
+from typing import Optional
 
 import kubernetes.client
-from mlrun_pipelines.mounts import mount_v3io, mount_v3iod
 
 import mlrun.errors
 from mlrun.config import config
+from mlrun.runtimes.mounts import mount_v3io, mount_v3iod
 
 from .kubejob import KubejobRuntime
 from .pod import KubeResourceSpec
@@ -179,7 +180,7 @@ class RemoteSparkRuntime(KubejobRuntime):
         skip_deployed=False,
         is_kfp=False,
         mlrun_version_specifier=None,
-        builder_env: dict = None,
+        builder_env: Optional[dict] = None,
         show_on_failure: bool = False,
         force_build: bool = False,
     ):

mlrun/runtimes/sparkjob/__init__.py

@@ -12,6 +12,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# flake8: noqa  - this is until we take care of the F401 violations with respect to __all__ & sphinx
-
 from .spark3job import Spark3Runtime

mlrun/runtimes/sparkjob/spark3job.py

@@ -14,13 +14,13 @@
 import typing
 
 import kubernetes.client
-from mlrun_pipelines.mounts import mount_v3io, mount_v3iod
 
 import mlrun.common.schemas.function
 import mlrun.errors
 import mlrun.k8s_utils
 import mlrun.runtimes.pod
 from mlrun.config import config
+from mlrun.runtimes.mounts import mount_v3io, mount_v3iod
 
 from ...execution import MLClientCtx
 from ...model import RunObject
@@ -405,8 +405,8 @@ class Spark3JobSpec(KubeResourceSpec):
     def _verify_and_set_requests(
         self,
         resources_field_name,
-        mem: str = None,
-        cpu: str = None,
+        mem: typing.Optional[str] = None,
+        cpu: typing.Optional[str] = None,
         patch: bool = False,
     ):
         # Spark operator uses JVM notation for memory, so we must verify it separately
@@ -774,7 +774,11 @@ class Spark3Runtime(KubejobRuntime):
                 exporter_jar="/spark/jars/jmx_prometheus_javaagent-0.16.1.jar",
             )
 
-    def with_cores(self, executor_cores: int = None, driver_cores: int = None):
+    def with_cores(
+        self,
+        executor_cores: typing.Optional[int] = None,
+        driver_cores: typing.Optional[int] = None,
+    ):
         """
         Allows to configure spark.executor.cores and spark.driver.cores parameters. The values must be integers
         greater than or equal to 1. If a parameter is not specified, it defaults to 1.
@@ -853,7 +857,7 @@ class Spark3Runtime(KubejobRuntime):
         skip_deployed=False,
         is_kfp=False,
         mlrun_version_specifier=None,
-        builder_env: dict = None,
+        builder_env: typing.Optional[dict] = None,
         show_on_failure: bool = False,
         force_build: bool = False,
     ):
@@ -944,7 +948,10 @@ class Spark3Runtime(KubejobRuntime):
         )
 
     def with_executor_requests(
-        self, mem: str = None, cpu: str = None, patch: bool = False
+        self,
+        mem: typing.Optional[str] = None,
+        cpu: typing.Optional[str] = None,
+        patch: bool = False,
     ):
         """
         set executor pod required cpu/memory/gpu resources
@@ -954,8 +961,8 @@ class Spark3Runtime(KubejobRuntime):
 
     def with_executor_limits(
         self,
-        cpu: str = None,
-        gpus: int = None,
+        cpu: typing.Optional[str] = None,
+        gpus: typing.Optional[int] = None,
         gpu_type: str = "nvidia.com/gpu",
         patch: bool = False,
     ):
@@ -970,7 +977,10 @@ class Spark3Runtime(KubejobRuntime):
         )
 
     def with_driver_requests(
-        self, mem: str = None, cpu: str = None, patch: bool = False
+        self,
+        mem: typing.Optional[str] = None,
+        cpu: typing.Optional[str] = None,
+        patch: bool = False,
     ):
         """
         set driver pod required cpu/memory/gpu resources
@@ -980,8 +990,8 @@ class Spark3Runtime(KubejobRuntime):
 
     def with_driver_limits(
         self,
-        cpu: str = None,
-        gpus: int = None,
+        cpu: typing.Optional[str] = None,
+        gpus: typing.Optional[int] = None,
         gpu_type: str = "nvidia.com/gpu",
         patch: bool = False,
     ):