fix: mutable default set argument in wait_for_job_status (#329)

Signed-off-by: Adithya <unknown0101200111@gmail.com>

Author: Adithya

kubeflow/optimizer/api/optimizer_client.py

@@ -18,7 +18,6 @@
 
 from kubeflow.common.types import KubernetesBackendConfig
 from kubeflow.optimizer.backends.kubernetes.backend import KubernetesBackend
-from kubeflow.optimizer.constants import constants
 from kubeflow.optimizer.types.algorithm_types import BaseAlgorithm
 from kubeflow.optimizer.types.optimization_types import (
     Objective,
@@ -182,7 +181,7 @@ def get_best_results(self, name: str) -> Result | None:
     def wait_for_job_status(
         self,
         name: str,
-        status: set[str] = {constants.OPTIMIZATION_JOB_COMPLETE},
+        status: set[str] | None = None,
         timeout: int = 3600,
         polling_interval: int = 2,
         callbacks: list[Callable[[OptimizationJob], None]] | None = None,

kubeflow/optimizer/backends/base.py

@@ -16,7 +16,6 @@
 from collections.abc import Callable, Iterator
 from typing import Any
 
-from kubeflow.optimizer.constants import constants
 from kubeflow.optimizer.types.algorithm_types import RandomSearch
 from kubeflow.optimizer.types.optimization_types import (
     Objective,
@@ -65,7 +64,7 @@ def get_best_results(self, name: str) -> Result | None:
     def wait_for_job_status(
         self,
         name: str,
-        status: set[str] = {constants.OPTIMIZATION_JOB_COMPLETE},
+        status: set[str] | None = None,
         timeout: int = 3600,
         polling_interval: int = 2,
         callbacks: list[Callable[[OptimizationJob], None]] | None = None,

kubeflow/optimizer/backends/kubernetes/backend.py

@@ -269,11 +269,14 @@ def get_best_results(self, name: str) -> Result | None:
     def wait_for_job_status(
         self,
         name: str,
-        status: set[str] = {constants.OPTIMIZATION_JOB_COMPLETE},
+        status: set[str] | None = None,
         timeout: int = 3600,
         polling_interval: int = 2,
         callbacks: list[Callable[[OptimizationJob], None]] | None = None,
     ) -> OptimizationJob:
+        if status is None:
+            status = {constants.OPTIMIZATION_JOB_COMPLETE}
+
         job_statuses = {
             constants.OPTIMIZATION_JOB_CREATED,
             constants.OPTIMIZATION_JOB_RUNNING,

kubeflow/trainer/api/trainer_client.py

@@ -233,7 +233,7 @@ def get_job_events(self, name: str) -> list[types.Event]:
     def wait_for_job_status(
         self,
         name: str,
-        status: set[str] = {constants.TRAINJOB_COMPLETE},
+        status: set[str] | None = None,
         timeout: int = 600,
         polling_interval: int = 2,
         callbacks: list[Callable[[types.TrainJob], None]] | None = None,

kubeflow/trainer/backends/base.py

@@ -75,7 +75,7 @@ def get_job_events(self, name: str) -> list[types.Event]:
     def wait_for_job_status(
         self,
         name: str,
-        status: set[str] = {constants.TRAINJOB_COMPLETE},
+        status: set[str] | None = None,
         timeout: int = 600,
         polling_interval: int = 2,
         callbacks: list[Callable[[types.TrainJob], None]] | None = None,

kubeflow/trainer/backends/container/backend.py

@@ -803,13 +803,16 @@ def get_job_events(self, name: str) -> list[types.Event]:
     def wait_for_job_status(
         self,
         name: str,
-        status: set[str] = {constants.TRAINJOB_COMPLETE},
+        status: set[str] | None = None,
         timeout: int = 600,
         polling_interval: int = 2,
         callbacks: list[Callable[[types.TrainJob], None]] | None = None,
     ) -> types.TrainJob:
         import time
 
+        if status is None:
+            status = {constants.TRAINJOB_COMPLETE}
+
         end = time.time() + timeout
         while time.time() < end:
             tj = self.get_job(name)

kubeflow/trainer/backends/kubernetes/backend.py

@@ -446,11 +446,14 @@ def get_job_logs(
     def wait_for_job_status(
         self,
         name: str,
-        status: set[str] = {constants.TRAINJOB_COMPLETE},
+        status: set[str] | None = None,
         timeout: int = 600,
         polling_interval: int = 2,
         callbacks: list[Callable[[types.TrainJob], None]] | None = None,
     ) -> types.TrainJob:
+        if status is None:
+            status = {constants.TRAINJOB_COMPLETE}
+
         job_statuses = {
             constants.TRAINJOB_CREATED,
             constants.TRAINJOB_RUNNING,

kubeflow/trainer/backends/localprocess/backend.py

@@ -215,11 +215,14 @@ def get_job_events(self, name: str) -> list[types.Event]:
     def wait_for_job_status(
         self,
         name: str,
-        status: set[str] = {constants.TRAINJOB_COMPLETE},
+        status: set[str] | None = None,
         timeout: int = 600,
         polling_interval: int = 2,
         callbacks: list[Callable[[types.TrainJob], None]] | None = None,
     ) -> types.TrainJob:
+        if status is None:
+            status = {constants.TRAINJOB_COMPLETE}
+
         # find first match or fallback
         _job = next((_job for _job in self.__local_jobs if _job.name == name), None)