ruff and isort

Author: b-yogesh

src/gaiaflow/cli/commands/minikube.py

@@ -12,11 +12,12 @@
 
 def load_imports():
     from gaiaflow.constants import BaseAction
-    from gaiaflow.managers.minikube_manager import (ExtendedAction,
-                                                    MinikubeManager)
-    from gaiaflow.managers.utils import (create_gaiaflow_context_path,
-                                         gaiaflow_path_exists_in_state,
-                                         parse_key_value_pairs)
+    from gaiaflow.managers.minikube_manager import ExtendedAction, MinikubeManager
+    from gaiaflow.managers.utils import (
+        create_gaiaflow_context_path,
+        gaiaflow_path_exists_in_state,
+        parse_key_value_pairs,
+    )
 
     return SimpleNamespace(
         BaseAction=BaseAction,
@@ -101,12 +102,14 @@ def restart(
     )
 
 
-@app.command(help="Containerize your package into a docker image inside the "
-                  "minikube cluster.")
+@app.command(
+    help="Containerize your package into a docker image inside the minikube cluster."
+)
 def dockerize(
     project_path: Path = typer.Option(..., "--path", "-p", help="Path to your project"),
-    image_name: str = typer.Option(DEFAULT_IMAGE_NAME, "--image-name", "-i",
-                                   help=("Name of your image.")),
+    image_name: str = typer.Option(
+        DEFAULT_IMAGE_NAME, "--image-name", "-i", help=("Name of your image.")
+    ),
 ):
     imports = load_imports()
     gaiaflow_path, user_project_path = imports.create_gaiaflow_context_path(
@@ -121,7 +124,7 @@ def dockerize(
         user_project_path=user_project_path,
         action=imports.ExtendedAction.DOCKERIZE,
         local=False,
-        image_name=image_name
+        image_name=image_name,
     )
 
 

src/gaiaflow/cli/commands/mlops.py

@@ -5,19 +5,21 @@
 import fsspec
 import typer
 
-from gaiaflow.constants import Service, DEFAULT_IMAGE_NAME
+from gaiaflow.constants import DEFAULT_IMAGE_NAME, Service
 
 app = typer.Typer()
 fs = fsspec.filesystem("file")
 
 
 def load_imports():
     from gaiaflow.constants import BaseAction, ExtendedAction
-    from gaiaflow.managers.mlops_manager import MlopsManager
     from gaiaflow.managers.minikube_manager import MinikubeManager
-    from gaiaflow.managers.utils import (create_gaiaflow_context_path,
-                                         gaiaflow_path_exists_in_state,
-                                         save_project_state)
+    from gaiaflow.managers.mlops_manager import MlopsManager
+    from gaiaflow.managers.utils import (
+        create_gaiaflow_context_path,
+        gaiaflow_path_exists_in_state,
+        save_project_state,
+    )
 
     return SimpleNamespace(
         BaseAction=BaseAction,
@@ -56,11 +58,17 @@ def start(
         False, "--docker-build", "-b", help="Force Docker image build"
     ),
     user_env_name: str = typer.Option(
-        None, "--env", "-e", help="Provide conda/mamba environment name for "
-                                 "Jupyter Lab to run. If not set, it will use the name from your environment.yml file."
+        None,
+        "--env",
+        "-e",
+        help="Provide conda/mamba environment name for "
+        "Jupyter Lab to run. If not set, it will use the name from your environment.yml file.",
     ),
     env_tool: "str" = typer.Option(
-        "mamba", "--env-tool", "-t", help="Which tool to use for running your Jupyter lab. Options: mamba, conda",
+        "mamba",
+        "--env-tool",
+        "-t",
+        help="Which tool to use for running your Jupyter lab. Options: mamba, conda",
     ),
 ):
     imports = load_imports()
@@ -242,12 +250,12 @@ def cleanup(
     )
 
 
-
 @app.command(help="Containerize your package into a docker image locally.")
 def dockerize(
     project_path: Path = typer.Option(..., "--path", "-p", help="Path to your project"),
-    image_name: str = typer.Option(DEFAULT_IMAGE_NAME, "--image-name", "-i",
-                                   help=("Name of your image.")),
+    image_name: str = typer.Option(
+        DEFAULT_IMAGE_NAME, "--image-name", "-i", help=("Name of your image.")
+    ),
 ):
     imports = load_imports()
     gaiaflow_path, user_project_path = imports.create_gaiaflow_context_path(
@@ -268,15 +276,18 @@ def dockerize(
         user_project_path=user_project_path,
         action=imports.ExtendedAction.DOCKERIZE,
         local=True,
-        image_name=image_name
+        image_name=image_name,
     )
 
-@app.command(help="Update the dependencies for the Airflow tasks. This command "
-                  "synchronizes the running container environments with the project's"
-                  "`environment.yml`. Make sure you have updated "
-                  "`environment.yml` before running"
-                  "this, as the container environments are updated based on "
-                  "its contents.")
+
+@app.command(
+    help="Update the dependencies for the Airflow tasks. This command "
+    "synchronizes the running container environments with the project's"
+    "`environment.yml`. Make sure you have updated "
+    "`environment.yml` before running"
+    "this, as the container environments are updated based on "
+    "its contents."
+)
 def update_deps(
     project_path: Path = typer.Option(..., "--path", "-p", help="Path to your project"),
 ):

src/gaiaflow/core/create_task.py

@@ -1,7 +1,11 @@
 from enum import Enum
 
-from .operators import (DevTaskOperator, DockerTaskOperator,
-                        ProdLocalTaskOperator, ProdTaskOperator)
+from .operators import (
+    DevTaskOperator,
+    DockerTaskOperator,
+    ProdLocalTaskOperator,
+    ProdTaskOperator,
+)
 
 
 class GaiaflowMode(Enum):

src/gaiaflow/core/operators.py

@@ -2,15 +2,15 @@
 import platform
 from datetime import datetime
 
-from airflow.providers.cncf.kubernetes.operators.pod import \
-    KubernetesPodOperator
+from airflow.providers.cncf.kubernetes.operators.pod import KubernetesPodOperator
 from airflow.providers.docker.operators.docker import DockerOperator
 from airflow.providers.standard.operators.python import ExternalPythonOperator
-from kubernetes.client import V1ResourceRequirements
 
-from gaiaflow.constants import (DEFAULT_MINIO_AWS_ACCESS_KEY_ID,
-                                DEFAULT_MINIO_AWS_SECRET_ACCESS_KEY,
-                                RESOURCE_PROFILES)
+from gaiaflow.constants import (
+    DEFAULT_MINIO_AWS_ACCESS_KEY_ID,
+    DEFAULT_MINIO_AWS_SECRET_ACCESS_KEY,
+    RESOURCE_PROFILES,
+)
 
 from .utils import build_env_from_secrets, inject_params_as_env_vars
 
@@ -128,8 +128,12 @@ def create_task(self):
 
         args, kwargs = self.resolve_args_kwargs()
         kwargs["params"] = dict(self.params)
-        op_kwargs = {"func_path": self.func_path, "args": args, "kwargs":
-            kwargs, "current_dir": current_dir}
+        op_kwargs = {
+            "func_path": self.func_path,
+            "args": args,
+            "kwargs": kwargs,
+            "current_dir": current_dir,
+        }
 
         def run_wrapper(**op_kwargs):
             import sys
@@ -206,17 +210,17 @@ def create_task(self):
         if profile is None:
             raise ValueError(f"Unknown resource profile: {profile_name}")
 
-        resources = V1ResourceRequirements(
-            requests={
-                "cpu": profile["request_cpu"],
-                "memory": profile["request_memory"],
-            },
-            limits={
-                "cpu": profile["limit_cpu"],
-                "memory": profile["limit_memory"],
-                # "gpu": profile.get["limit_gpu"],
-            },
-        )
+        # resources = V1ResourceRequirements(
+        #     requests={
+        #         "cpu": profile["request_cpu"],
+        #         "memory": profile["request_memory"],
+        #     },
+        #     limits={
+        #         "cpu": profile["limit_cpu"],
+        #         "memory": profile["limit_memory"],
+        #         # "gpu": profile.get["limit_gpu"],
+        #     },
+        # )
 
         return KubernetesPodOperator(
             task_id=self.task_id,

src/gaiaflow/core/runner.py

@@ -9,6 +9,7 @@
 import pickle
 from typing import Any
 
+
 def run(
     func_path: str | None = None,
     args: list | None = None,
@@ -39,30 +40,34 @@ def _extract_params_from_env(prefix="PARAMS_") -> dict[str, str]:
         if k.startswith(prefix)
     }
 
-def _resolve_inputs(func_path: str, args: list[Any], kwargs: dict[Any],
-                   mode: str):
+
+def _resolve_inputs(func_path: str, args: list[Any], kwargs: dict[Any], mode: str):
     if mode == "dev":
         return func_path, args or [], kwargs or {}
-    else: # all other modes (dev_docker, prod_local and prod)
+    else:  # all other modes (dev_docker, prod_local and prod)
         func_path = os.environ.get("FUNC_PATH", func_path)
         args = json.loads(os.environ.get("FUNC_ARGS", "[]"))
         kwargs = json.loads(os.environ.get("FUNC_KWARGS", "{}"))
         kwargs["params"] = _extract_params_from_env()
         return func_path, args, kwargs
 
+
 def _import_function(func_path: str):
     import importlib
+
     module_path, func_name = func_path.rsplit(":", 1)
     module = importlib.import_module(module_path)
     return getattr(module, func_name)
 
+
 def _write_result(result, mode):
     if mode == "prod" or mode == "prod_local":
         _write_xcom_result(result)
     if mode == "dev_docker":
         with open("/tmp/script.out", "wb+") as tmp:
             pickle.dump(result, tmp)
 
+
 def _write_xcom_result(result: Any) -> None:
     try:
         xcom_dir = "/airflow/xcom"

src/gaiaflow/core/utils.py

@@ -28,8 +28,7 @@ def docker_network_gateway() -> str | None:
             if "Gateway" in line:
                 match = re.search(r'"Gateway"\s*:\s*"([^"]+)"', line)
                 if match:
-                    print(f"Docker network Gateway for Minikube is - "
-                          f"{match.group(1)}")
+                    print(f"Docker network Gateway for Minikube is - {match.group(1)}")
                     return match.group(1)
         print("Is your minikube cluster running? Please run and try again.")
         return None
@@ -41,5 +40,6 @@ def docker_network_gateway() -> str | None:
         print("Docker command not found. Is Docker installed and in your PATH?")
         return None
 
+
 if __name__ == "__main__":
-    docker_network_gateway()
+    docker_network_gateway()