diff --git a/CHANGES.md b/CHANGES.md index 4774652..5cd9a7d 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -1,3 +1,9 @@ +## Version 0.0.4 (in development) + +* Added tests +* Internal refactoring + + ## Version 0.0.3 * **Bug fix** - `ExternalPythonOperator` does not need Airflow in external environment now. diff --git a/docs/index.md b/docs/index.md index b0b2abc..a295a2a 100644 --- a/docs/index.md +++ b/docs/index.md @@ -215,7 +215,7 @@ Quick rule of thumb: - Use `prod_local` for testing end-to-end workflows on production-like settngs. - Use `prod` for production pipelines in the real cluster. -## User workflow: +## User workflow A typical user workflow could look like this: ```mermaid diff --git a/pixi.lock b/pixi.lock index 0cf2fe8..fa3db2e 100644 --- a/pixi.lock +++ b/pixi.lock @@ -2382,8 +2382,8 @@ packages: timestamp: 1741505873840 - pypi: ./ name: gaiaflow - version: 0.0.3.dev0 - sha256: 64847b19a8e8c898e9c712054f2aecbdc022c3a2a8a6584fa7995cb5f7d69067 + version: 0.0.4.dev0 + sha256: 5e3ceaa3916b7d3961eb01d16f9b1ee62a9c146a0ae379f3aab37f3fe54b1b78 requires_dist: - typer>=0.16.0,<0.17 - fsspec>=2025.7.0,<2026 diff --git a/pyproject.toml b/pyproject.toml index 67d417b..4060124 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "gaiaflow" requires-python = ">= 3.11" -version = "0.0.3" +version = "0.0.4dev0" description = "Local-first MLOps infrastructure python tool that simplifies the process of building, testing, and deploying ML workflows." authors = [{name = "Yogesh Kumar Baljeet Singh", email = "yogesh.baljeetsingh@brockmann-consult.de"}] dependencies = [ diff --git a/src/gaiaflow/cli/commands/minikube.py b/src/gaiaflow/cli/commands/minikube.py index 687b803..755c41f 100644 --- a/src/gaiaflow/cli/commands/minikube.py +++ b/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, ) diff --git a/src/gaiaflow/cli/commands/mlops.py b/src/gaiaflow/cli/commands/mlops.py index a80447e..5962a8b 100644 --- a/src/gaiaflow/cli/commands/mlops.py +++ b/src/gaiaflow/cli/commands/mlops.py @@ -5,7 +5,7 @@ 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") @@ -13,11 +13,13 @@ 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"), ): diff --git a/src/gaiaflow/core/create_task.py b/src/gaiaflow/core/create_task.py index f0a3078..7dd4f7a 100755 --- a/src/gaiaflow/core/create_task.py +++ b/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): diff --git a/src/gaiaflow/core/operators.py b/src/gaiaflow/core/operators.py index 9b673a2..f481865 100644 --- a/src/gaiaflow/core/operators.py +++ b/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 @@ -122,16 +122,35 @@ def create_func_env_vars(self): class DevTaskOperator(BaseTaskOperator): def create_task(self): - from gaiaflow.core.runner import run + import os + + current_dir = os.path.dirname(os.path.abspath(__file__)) args, kwargs = self.resolve_args_kwargs() kwargs["params"] = dict(self.params) - op_kwargs = {"func_path": self.func_path, "args": args, "kwargs": kwargs} + op_kwargs = { + "func_path": self.func_path, + "args": args, + "kwargs": kwargs, + "current_dir": current_dir, + } + + def run_wrapper(**op_kwargs): + import sys + + sys.path.append(op_kwargs.get("current_dir", "")) + from runner import run + + return run( + func_path=op_kwargs.get("func_path"), + args=op_kwargs.get("args"), + kwargs=op_kwargs.get("kwargs"), + ) return ExternalPythonOperator( task_id=self.task_id, python="/home/airflow/.local/share/mamba/envs/default_user_env/bin/python", - python_callable=run, + python_callable=run_wrapper, op_kwargs=op_kwargs, do_xcom_push=True, retries=self.retries, @@ -191,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, @@ -216,7 +235,7 @@ def create_task(self): do_xcom_push=True, retries=self.retries, params=self.params, - container_resources=resources, + # container_resources=resources, ) diff --git a/src/gaiaflow/core/runner.py b/src/gaiaflow/core/runner.py index e5f1b74..9ea9b19 100644 --- a/src/gaiaflow/core/runner.py +++ b/src/gaiaflow/core/runner.py @@ -10,6 +10,29 @@ from typing import Any +def run( + func_path: str | None = None, + args: list | None = None, + kwargs: dict[str, Any] | None = None, +) -> dict[str, str]: + mode = os.environ.get("MODE", "dev") + print(f"## Runner running in {mode} mode ##") + func_path, args, kwargs = _resolve_inputs(func_path, args, kwargs, mode) + + if not func_path: + raise ValueError("func_path must be provided") + + func = _import_function(func_path) + + print(f"Running {func_path} with args: {args} and kwargs :{kwargs}") + result = func(*args, **kwargs) + print("Function result:", result) + + _write_result(result, mode) + + return result + + def _extract_params_from_env(prefix="PARAMS_") -> dict[str, str]: return { k[len(prefix) :].lower(): v @@ -18,47 +41,32 @@ def _extract_params_from_env(prefix="PARAMS_") -> dict[str, str]: } -def run( - func_path: str | None = None, - args: list | None = None, - kwargs: dict[str, Any] | None = None, -) -> dict[str, str]: - mode = os.environ.get("MODE", "dev") - print(f"## Runner running in {mode} mode ##") +def _resolve_inputs(func_path: str, args: list[Any], kwargs: dict[Any], mode: str): if mode == "dev": - print("args", args) - print("kwargs", kwargs) - else: - func_path = os.environ.get("FUNC_PATH", "") + return func_path, args or [], kwargs or {} + 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", "{}")) - params: dict = _extract_params_from_env() - kwargs["params"] = params - print("args", args) - print("kwargs", kwargs) + kwargs["params"] = _extract_params_from_env() + return func_path, args, kwargs - if not func_path: - raise ValueError("func_path must be provided") - module_path, func_name = func_path.rsplit(":", 1) +def _import_function(func_path: str): import importlib + module_path, func_name = func_path.rsplit(":", 1) module = importlib.import_module(module_path) - func = getattr(module, func_name) + return getattr(module, func_name) - print(f"Running {func_path} with args: {args} and kwargs :{kwargs}") - result = func(*args, **kwargs) - print("Function result:", result) + +def _write_result(result, mode): if mode == "prod" or mode == "prod_local": - # This is needed when we use KubernetesPodOperator and want to - # share information via XCOM. _write_xcom_result(result) if mode == "dev_docker": with open("/tmp/script.out", "wb+") as tmp: pickle.dump(result, tmp) - return result - def _write_xcom_result(result: Any) -> None: try: @@ -68,12 +76,6 @@ def _write_xcom_result(result: Any) -> None: with open(f"{xcom_dir}/return.json", "w") as f: json.dump(result, f) - path = "/airflow/xcom/return.json" - print("[DEBUG] File exists:", os.path.exists(path)) - print("[DEBUG] File size:", os.path.getsize(path)) - with open(path, "r") as f: - print("[DEBUG] File contents:", f.read()) - print("Result written to XCom successfully") except Exception as e: print(f"Failed to write XCom result: {e}") diff --git a/src/gaiaflow/core/utils.py b/src/gaiaflow/core/utils.py index a053b44..0093dae 100644 --- a/src/gaiaflow/core/utils.py +++ b/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 @@ -40,3 +39,7 @@ def docker_network_gateway() -> str | None: except FileNotFoundError: print("Docker command not found. Is Docker installed and in your PATH?") return None + + +if __name__ == "__main__": + docker_network_gateway() diff --git a/src/gaiaflow/managers/base_manager.py b/src/gaiaflow/managers/base_manager.py index 54ab57e..263ba08 100644 --- a/src/gaiaflow/managers/base_manager.py +++ b/src/gaiaflow/managers/base_manager.py @@ -17,7 +17,7 @@ def __init__( valid_actions = self._get_valid_actions() if action not in valid_actions: valid_names = sorted([a.name for a in valid_actions]) - raise ValueError(f"Invalid action '{action.name}'. Valid: {valid_names}") + raise ValueError(f"Invalid action '{action}'. Valid: {valid_names}") self.gaiaflow_path = gaiaflow_path self.user_project_path = user_project_path diff --git a/src/gaiaflow/managers/minikube_manager.py b/src/gaiaflow/managers/minikube_manager.py index 58d2942..1182a36 100644 --- a/src/gaiaflow/managers/minikube_manager.py +++ b/src/gaiaflow/managers/minikube_manager.py @@ -8,13 +8,24 @@ import yaml -from gaiaflow.constants import (AIRFLOW_SERVICES, MINIO_SERVICES, - MLFLOW_SERVICES, Action, BaseAction, - ExtendedAction) +from gaiaflow.constants import ( + AIRFLOW_SERVICES, + MINIO_SERVICES, + MLFLOW_SERVICES, + Action, + BaseAction, + ExtendedAction, +) from gaiaflow.managers.base_manager import BaseGaiaflowManager from gaiaflow.managers.mlops_manager import MlopsManager -from gaiaflow.managers.utils import (find_python_packages, log_error, log_info, - run, set_permissions, is_wsl) +from gaiaflow.managers.utils import ( + find_python_packages, + is_wsl, + log_error, + log_info, + run, + set_permissions, +) @contextmanager @@ -28,6 +39,255 @@ def temporary_copy(src: Path, dest: Path): dest.unlink() +class MinikubeHelper: + def __init__(self, profile: str = "airflow"): + self.profile = profile + + def is_running(self) -> bool: + result = subprocess.run( + ["minikube", "status", "--profile", self.profile], + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + ) + return b"Running" in result.stdout + + def start(self): + if self.is_running(): + log_info(f"Minikube cluster [{self.profile}] is already running.") + return + + log_info(f"Starting Minikube cluster [{self.profile}]...") + cmd = [ + "minikube", + "start", + "--profile", + self.profile, + "--driver=docker", + "--cpus=4", + "--memory=4g", + ] + if is_wsl(): + cmd.append("--extra-config=kubelet.cgroup-driver=cgroupfs") + + try: + run(cmd, f"Error starting minikube profile [{self.profile}]") + except subprocess.CalledProcessError: + log_info("Retrying after cleanup...") + self.cleanup() + run(cmd, f"Error starting minikube profile [{self.profile}]") + + def stop(self): + log_info(f"Stopping minikube profile [{self.profile}]...") + run( + ["minikube", "stop", "--profile", self.profile], + f"Error stopping minikube profile [{self.profile}]", + ) + + def cleanup(self): + log_info(f"Deleting minikube profile: {self.profile}") + run( + ["minikube", "delete", "--profile", self.profile], + f"Error deleting minikube profile [{self.profile}]", + ) + + def run_cmd(self, args: list[str], **kwargs): + full_cmd = ["minikube", "-p", self.profile] + args + return subprocess.run(full_cmd, **kwargs) + + +class DockerHelper: + def __init__( + self, + image_name: str, + project_path: Path, + local: bool, + minikube_helper: MinikubeHelper, + ): + self.image_name = image_name + self.project_path = project_path + self.local = local + self.minikube_helper = minikube_helper + + def build_image(self, dockerfile_path: Path): + if not dockerfile_path.exists(): + log_error(f"Dockerfile not found at {dockerfile_path}") + return + + log_info(f"Updating Dockerfile at {dockerfile_path}") + self._update_dockerfile(dockerfile_path) + + runner_src = Path(__file__).parent.parent.resolve() / "core" / "runner.py" + runner_dest = self.project_path / "runner.py" + + with temporary_copy(runner_src, runner_dest): + if self.local: + self._build_local(dockerfile_path) + else: + self._build_minikube(dockerfile_path) + + def _update_dockerfile(self, dockerfile_path: Path): + DockerHelper._add_copy_statements_to_dockerfile( + str(dockerfile_path), find_python_packages(self.project_path) + ) + + def _build_local(self, dockerfile_path: Path): + log_info(f"Building Docker image [{self.image_name}] locally") + run( + [ + "docker", + "build", + "-t", + self.image_name, + "-f", + dockerfile_path, + self.project_path, + ], + "Error building Docker image locally", + ) + set_permissions("/var/run/docker.sock", 0o666) + + def _build_minikube(self, dockerfile_path: Path): + log_info(f"Building Docker image [{self.image_name}] in Minikube context") + result = self.minikube_helper.run_cmd( + ["docker-env", "--shell", "bash"], stdout=subprocess.PIPE, check=True + ) + env = self._parse_minikube_env(result.stdout.decode()) + run( + [ + "docker", + "build", + "-t", + self.image_name, + "-f", + dockerfile_path, + self.project_path, + ], + "Error building Docker image inside Minikube", + env=env, + ) + + @staticmethod + def _parse_minikube_env(output: str) -> dict: + env = os.environ.copy() + for line in output.splitlines(): + if line.startswith("export "): + try: + key, value = line.replace("export ", "").split("=", 1) + env[key.strip()] = value.strip('"') + except ValueError: + continue + return env + + @staticmethod + def _add_copy_statements_to_dockerfile( + dockerfile_path: str, local_packages: list[str] + ): + with open(dockerfile_path, "r") as f: + lines = f.readlines() + + env_index = next( + (i for i, line in enumerate(lines) if line.strip().startswith("ENV")), + None, + ) + + if env_index is None: + raise ValueError("No ENV found in Dockerfile.") + + entrypoint_index = next( + ( + i + for i, line in enumerate(lines) + if line.strip().startswith("ENTRYPOINT") + ), + None, + ) + + if entrypoint_index is None: + raise ValueError("No ENTRYPOINT found in Dockerfile.") + + copy_lines = [f"COPY {pkg} ./{pkg}\n" for pkg in local_packages] + copy_lines.append("COPY runner.py ./runner.py\n") + + updated_lines = ( + lines[: env_index + 1] + + copy_lines # + + lines[entrypoint_index:] + ) + with open(dockerfile_path, "w") as f: + f.writelines(updated_lines) + + print("Dockerfile updated with COPY statements.") + + +class KubeConfigHelper: + def __init__(self, gaiaflow_path: Path, os_type: str): + self.gaiaflow_path = gaiaflow_path + self.os_type = os_type + + def create_inline(self): + kube_config = Path.home() / ".kube" / "config" + backup_config = kube_config.with_suffix(".backup") + + self._backup_kube_config(kube_config, backup_config) + self._patch_kube_config(kube_config) + self._write_inline(kube_config) + + if (self.os_type == "windows" or is_wsl()) and backup_config.exists(): + shutil.copy(backup_config, kube_config) + backup_config.unlink() + log_info("Reverted kube config to original state.") + + def _backup_kube_config(self, kube_config: Path, backup_config: Path): + if kube_config.exists(): + with open(kube_config, "r") as f: + config_data = yaml.safe_load(f) + with open(backup_config, "w") as f: + yaml.dump(config_data, f) + + def _patch_kube_config(self, kube_config: Path): + if not kube_config.exists(): + return + + with open(kube_config, "r") as f: + config_data = yaml.safe_load(f) + + for cluster in config_data.get("clusters", []): + cluster_info = cluster.get("cluster", {}) + if self.os_type == "windows": + server = cluster_info.get("server", "") + if "127.0.0.1" in server or "localhost" in server: + cluster_info["server"] = server.replace( + "127.0.0.1", "host.docker.internal" + ).replace("localhost", "host.docker.internal") + cluster_info["insecure-skip-tls-verify"] = True + elif is_wsl(): + cluster_info["server"] = "https://192.168.49.2:8443" + cluster_info["insecure-skip-tls-verify"] = True + + with open(kube_config, "w") as f: + yaml.dump(config_data, f) + + def _write_inline(self, kube_config: Path): + filename = self.gaiaflow_path / "docker_stuff" / "kube_config_inline" + log_info("Creating kube config inline file...") + with open(filename, "w") as f: + subprocess.call( + [ + "minikube", + "kubectl", + "--", + "config", + "view", + "--flatten", + "--minify", + "--raw", + ], + cwd=self.gaiaflow_path / "docker_stuff", + stdout=f, + ) + log_info(f"Created kube config inline file {filename}") + + class MinikubeManager(BaseGaiaflowManager): def __init__( self, @@ -50,6 +310,17 @@ def __init__( self.local = local self.image_name = image_name + self.minikube_helper = MinikubeHelper() + self.docker_helper = DockerHelper( + image_name=image_name, + project_path=user_project_path, + local=local, + minikube_helper=self.minikube_helper, + ) + self.kube_helper = KubeConfigHelper( + gaiaflow_path=gaiaflow_path, os_type=self.os_type + ) + super().__init__( gaiaflow_path=gaiaflow_path, user_project_path=user_project_path, @@ -59,13 +330,11 @@ def __init__( ) def _get_valid_actions(self) -> Set[Action]: - base_actions = super()._get_valid_actions() - extra_actions = { + return super()._get_valid_actions() | { ExtendedAction.DOCKERIZE, ExtendedAction.CREATE_CONFIG, ExtendedAction.CREATE_SECRET, } - return base_actions | extra_actions @classmethod def run(cls, **kwargs): @@ -73,7 +342,7 @@ def run(cls, **kwargs): if action is None: raise ValueError("Missing required argument 'action'") - manager = MinikubeManager(**kwargs) + manager = cls(**kwargs) action_map = { BaseAction.START: manager.start, @@ -82,57 +351,24 @@ def run(cls, **kwargs): BaseAction.CLEANUP: manager.cleanup, ExtendedAction.DOCKERIZE: manager.build_docker_image, ExtendedAction.CREATE_CONFIG: manager.create_kube_config_inline, - ExtendedAction.CREATE_SECRET: manager.create_secrets, + ExtendedAction.CREATE_SECRET: lambda: manager.create_secrets( + kwargs["secret_name"], kwargs["secret_data"] + ), } try: - action_method = action_map[action] + action_map[action]() except KeyError: raise ValueError(f"Unknown action: {action}") - if action == ExtendedAction.CREATE_SECRET: - action_method(kwargs["secret_name"], kwargs["secret_data"]) - else: - action_method() - - def start(self): - if self.force_new: - self.cleanup() - MlopsManager.run(gaiaflow_path=self.gaiaflow_path, user_project_path=self.user_project_path, action=BaseAction.STOP) - log_info(f"Checking Minikube cluster [{self.minikube_profile}] status...") - result = subprocess.run( - ["minikube", "status", "--profile", self.minikube_profile], - stdout=subprocess.PIPE, - stderr=subprocess.DEVNULL, + def _stop_mlops(self): + MlopsManager.run( + gaiaflow_path=self.gaiaflow_path, + user_project_path=self.user_project_path, + action=BaseAction.STOP, ) - if b"Running" in result.stdout: - log_info(f"Minikube cluster [{self.minikube_profile}] is already running.") - else: - log_info( - f"Minikube cluster [{self.minikube_profile}] is not running. Starting..." - ) - try: - cmd = [ - "minikube", - "start", - "--profile", - self.minikube_profile, - "--driver=docker", - "--cpus=4", - "--memory=4g", - ] - if is_wsl(): - cmd.append("--extra-config=kubelet.cgroup-driver=cgroupfs") - run( - cmd, - f"Error starting minikube profile [{self.minikube_profile}]", - ) - except subprocess.CalledProcessError: - log_info("Cleaning up and starting again...") - self.cleanup() - self.start() - self.create_kube_config_inline() + def _start_mlops(self): MlopsManager.run( gaiaflow_path=self.gaiaflow_path, user_project_path=self.user_project_path, @@ -141,197 +377,25 @@ def start(self): force_new=self.force_new, ) + def start(self): + if self.force_new: + self.cleanup() + self._stop_mlops() + self.minikube_helper.start() + self.create_kube_config_inline() + self._start_mlops() + def stop(self): - log_info(f"Stopping minikube profile [{self.minikube_profile}]...") - try: - run( - ["minikube", "stop", "--profile", self.minikube_profile], - f"Error stopping minikube profile [{self.minikube_profile}]", - ) - log_info(f"Stopped minikube profile [{self.minikube_profile}]") - except Exception as e: - log_info(str(e)) + self.minikube_helper.stop() def create_kube_config_inline(self): - kube_config = Path.home() / ".kube" / "config" - backup_config = kube_config.with_suffix(".backup") - filename = f"{self.gaiaflow_path / 'docker_stuff'}/kube_config_inline" - - if kube_config.exists(): - with open(kube_config, "r") as f: - config_data = yaml.safe_load(f) - - with open(backup_config, "w") as f: - yaml.dump(config_data, f) - - for cluster in config_data.get("clusters", []): - if self.os_type == "windows" and kube_config.exists(): - log_info("Detected Windows: patching kube config with host.docker.internal") - server = cluster.get("cluster", {}).get("server", "") - if "127.0.0.1" in server or "localhost" in server: - cluster["cluster"]["server"] = server.replace( - "127.0.0.1", "host.docker.internal" - ).replace("localhost", "host.docker.internal") - - - elif is_wsl(): - log_info("Detected WSL: patching kube config with minikube ip") - # ip = subprocess.check_output(["minikube", "ip", "-p", "airflow"], text=True).strip() - # cluster["cluster"]["server"] = f"https://{ip}:8443" - cluster["cluster"]["server"] = "https://192.168.49.2:8443" - cluster["cluster"]["insecure-skip-tls-verify"] = True - - with open(kube_config, "w") as f: - yaml.dump(config_data, f) - - log_info("Creating kube config inline file...") - with open(filename, "w") as f: - subprocess.call( - [ - "minikube", - "kubectl", - "--", - "config", - "view", - "--flatten", - "--minify", - "--raw", - ], - cwd=self.gaiaflow_path / "docker_stuff", - stdout=f, - ) - - log_info(f"Created kube config inline file {filename}") - - if self.os_type == "windows": - log_info( - f"Adding insecure-skip-tls-verfiy for local setup in kube config inline file {filename}" - ) - with open(filename, "r") as f: - kube_config_data = yaml.safe_load(f) - - - for cluster in kube_config_data.get("clusters", []): - cluster_data = cluster.get("cluster", {}) - if "insecure-skip-tls-verify" not in cluster_data: - cluster_data["insecure-skip-tls-verify"] = True - - log_info(f"Saving kube config inline file {filename}") - with open(filename, "w") as f: - yaml.safe_dump(kube_config_data, f, default_flow_style=False) - - if (self.os_type == "windows" or is_wsl()) and backup_config.exists(): - shutil.copy(backup_config, kube_config) - backup_config.unlink() - log_info("Reverted kube config to original state.") - - @staticmethod - def _add_copy_statements_to_dockerfile( - dockerfile_path: str, local_packages: list[str] - ): - with open(dockerfile_path, "r") as f: - lines = f.readlines() - - env_index = next( - (i for i, line in enumerate(lines) if line.strip().startswith("ENV")), - None, - ) - - if env_index is None: - raise ValueError("No ENV found in Dockerfile.") - - entrypoint_index = next( - ( - i - for i, line in enumerate(lines) - if line.strip().startswith("ENTRYPOINT") - ), - None, - ) - - if entrypoint_index is None: - raise ValueError("No ENTRYPOINT found in Dockerfile.") - - copy_lines = [f"COPY {pkg} ./{pkg}\n" for pkg in local_packages] - copy_lines.append("COPY runner.py ./runner.py\n") - - updated_lines = ( - lines[: env_index + 1] - + copy_lines # - + lines[entrypoint_index:] - ) - with open(dockerfile_path, "w") as f: - f.writelines(updated_lines) - - print("Dockerfile updated with COPY statements.") + self.kube_helper.create_inline() def build_docker_image(self): - dockerfile_path = self.gaiaflow_path / "docker_stuff" / "user-package" / "Dockerfile" - if not (dockerfile_path.exists()): - log_error(f"Dockerfile not found at {dockerfile_path}") - return - - log_info(f"Updating dockerfile at {dockerfile_path}") - MinikubeManager._add_copy_statements_to_dockerfile( - dockerfile_path, find_python_packages(self.user_project_path) + dockerfile_path = ( + self.gaiaflow_path / "docker_stuff" / "user-package" / "Dockerfile" ) - runner_src = Path(__file__).parent.parent.resolve() / "core" / "runner.py" - runner_dest = self.user_project_path / "runner.py" - - with temporary_copy(runner_src, runner_dest): - if self.local: - log_info(f"Building Docker image [{self.image_name}] locally") - run( - [ - "docker", - "build", - "-t", - self.image_name, - "-f", - dockerfile_path, - self.user_project_path, - ], - "Error building docker image.", - ) - # TODO: For windows? - set_permissions("/var/run/docker.sock", 0o666) - else: - log_info( - f"Building Docker image [{self.image_name}] in minikube context" - ) - result = subprocess.run( - [ - "minikube", - "-p", - self.minikube_profile, - "docker-env", - "--shell", - "bash", - ], - stdout=subprocess.PIPE, - check=True, - ) - env = os.environ.copy() - for line in result.stdout.decode().splitlines(): - if line.startswith("export "): - try: - key, value = line.replace("export ", "").split("=", 1) - env[key.strip()] = value.strip('"') - except ValueError: - continue - run( - [ - "docker", - "build", - "-t", - self.image_name, - "-f", - dockerfile_path, - self.user_project_path, - ], - "Error building docker image inside minikube cluster.", - env=env, - ) + self.docker_helper.build_image(dockerfile_path) def create_secrets(self, secret_name: str, secret_data: dict[str, Any]): log_info(f"Checking if secret [{secret_name}] exists...") diff --git a/src/gaiaflow/managers/mlflow_model_manager.py b/src/gaiaflow/managers/mlflow_model_manager.py index 1206b6b..c14d2e5 100644 --- a/src/gaiaflow/managers/mlflow_model_manager.py +++ b/src/gaiaflow/managers/mlflow_model_manager.py @@ -1,204 +1,204 @@ -import os -import subprocess -import uuid - -import mlflow - -import docker -from gaiaflow.constants import (DEFAULT_MINIO_AWS_ACCESS_KEY_ID, - DEFAULT_MINIO_AWS_SECRET_ACCESS_KEY, - DEFAULT_MLFLOW_TRACKING_URI, BaseAction) -from gaiaflow.managers.base_manager import BaseGaiaflowManager - - -class MlflowModelManager(BaseGaiaflowManager): - def __init__( - self, - registry_uri=None, - action: BaseAction = None, - force_new: bool = False, - prune: bool = False, - local: bool = False, - **kwargs, - ): - self.docker_client = docker.from_env() - self.registry_uri = registry_uri.rstrip("/") if registry_uri else None - self.local = local - - # TODO: not optimal. Use ContextManager to restore the env vars - # Move them to the actual methods where they are needed. - os.environ["MLFLOW_TRACKING_URI"] = ( - os.getenv("MLFLOW_TRACKING_URI") or "https://localhost:5000" - ) - os.environ["MLFLOW_S3_ENDPOINT_URL"] = ( - os.getenv("MLFLOW_S3_ENDPOINT_URL") or "https://localhost:9000" - ) - os.environ["AWS_ACCESS_KEY_ID"] = ( - os.getenv("AWS_ACCESS_KEY_ID") or DEFAULT_MINIO_AWS_ACCESS_KEY_ID - ) - os.environ["AWS_SECRET_ACCESS_KEY"] = ( - os.getenv("AWS_SECRET_ACCESS_KEY") or DEFAULT_MINIO_AWS_SECRET_ACCESS_KEY - ) - - mlflow_tracking_uri = ( - os.getenv("MLFLOW_TRACKING_URI") or DEFAULT_MLFLOW_TRACKING_URI - ) - mlflow.set_tracking_uri(mlflow_tracking_uri) - - # This Manager does not need the paths. But it expects that either - # the model exists in S3 if s3_uri is used or a live instance of MLFlow - # running if run_id is used - super().__init__( - gaiaflow_path="", - user_project_path="", - action=action, - force_new=force_new, - prune=prune, - **kwargs, - ) - - def _get_model_uri(self, model_uri=None, run_id=None): - if model_uri: - return model_uri - elif run_id: - return f"runs:/{run_id}/model" - else: - raise ValueError("Either model_uri or run_id must be provided.") - - # TODO: Use this instead - # @classmethod - # def run(cls, action, **kwargs): - # params = kwargs.get("params", {}) - # model_uri = params.get("model_uri") - # run_id = params.get("run_id") - # image_name = params.get("image_name") - # enable_mlserver = params.get("enable_mlserver", True) - # manag = MlflowModelManager(model_uri, run_id, image_name, - # enable_mlserver) - # environment = create_env() - # with set_env_cm(environment): - # if action == BaseAction.START: - # manag.start() - - def start(self, **kwargs): - print("kwargs inside start, **kwargs", kwargs) - params = kwargs.get("params", {}) - - model_uri = params.get("model_uri") - run_id = params.get("run_id") - image_name = params.get("image_name") - enable_mlserver = params.get("enable_mlserver", True) - - self.build_model_docker_image(model_uri, run_id, image_name, enable_mlserver) - self.run_container(image_name) - - def stop(self, **kwargs): - print("kwargs inside stop, **kwargs", kwargs) - params = kwargs.get("params", {}) - - container_id_or_name = params.get("container_id_or_name") - self.stop_and_remove_container(container_id_or_name) - - def cleanup(self, **kwargs): - print("kwargs inside cleanup, **kwargs", kwargs) - params = kwargs.get("params", {}) - - container_id_or_name = params.get("container_id_or_name") - purge = params.get("purge") - self.stop_and_remove_container(container_id_or_name, purge) - - def build_model_docker_image( - self, model_uri=None, run_id=None, image_name=None, enable_mlserver=True - ): - model_uri_final = self._get_model_uri(model_uri, run_id) - - unique_tag = str(uuid.uuid4())[:8] - image_name = image_name or f"mlflow-model:{unique_tag}" - - full_image_name = ( - f"{self.registry_uri}/{image_name}" if self.registry_uri else image_name - ) - - print(f"Building Docker image from {model_uri_final} as {full_image_name}...") - - mlflow.models.build_docker( - model_uri=model_uri_final, - name=full_image_name, - enable_mlserver=enable_mlserver, - ) - - return full_image_name - - def run_container(self, image_name, port=8246): - container = self.docker_client.containers.run( - image_name, - detach=True, - ports={"8080/tcp": port}, - name=f"mlflow_model_{uuid.uuid4().hex[:6]}", - ) - print(f"Container started: {container.name} on port {port}") - return container - - def stop_and_remove_container(self, container_id_or_name, purge=False): - try: - container = self.docker_client.containers.get(container_id_or_name) - print(f"Stopping container: {container.name}") - container.stop() - container.remove() - print("Container stopped and removed.") - if purge: - image_name = container.image.tags[0] - self.docker_client.images.remove(image=image_name, force=True) - print(f"Image {image_name} removed.") - except docker.context.api.errors.NotFound: - print("Container not found.") - - def push_image_to_registry(self, image_name): - if not self.registry_uri: - raise ValueError("No registry_uri provided.") - print(f"Pushing image {image_name} to registry {self.registry_uri}...") - subprocess.run(["docker", "push", image_name], check=True) - print("Image pushed.") - - def generate_k8s_yaml(self): - print(" Generating Kubernetes deployment YAML...") - raise NotImplementedError() - - def deploy_to_kubernetes(self, k8s_yaml_path): - print(f"Deploying to Kubernetes from {k8s_yaml_path}...") - subprocess.run(["kubectl", "apply", "-f", k8s_yaml_path], check=True) - print("Deployment applied.") - - -# -# @contextmanager -# def set_env_cm(**new_env: str | None) -> Generator[dict[str, str | None]]: -# """Run the code in the block with a new environment `new_env`.""" -# restore_env = set_env(**new_env) -# try: -# yield new_env -# finally: -# restore_env() -# -# -# def set_env(**new_env: str | None) -> Callable[[], None]: -# """ -# Set the new environment in `new_env` and return a no-arg -# function to restore the old environment. -# """ -# old_env = {k: os.environ.get(k) for k in new_env.keys()} -# -# def restore_env(): -# for ko, vo in old_env.items(): -# if vo is not None: -# os.environ[ko] = vo -# elif ko in os.environ: -# del os.environ[ko] -# -# for kn, vn in new_env.items(): -# if vn is not None: -# os.environ[kn] = vn -# elif kn in os.environ: -# del os.environ[kn] -# -# return restore_env +# import os +# import subprocess +# import uuid +# +# import mlflow +# +# import docker +# from gaiaflow.constants import (DEFAULT_MINIO_AWS_ACCESS_KEY_ID, +# DEFAULT_MINIO_AWS_SECRET_ACCESS_KEY, +# DEFAULT_MLFLOW_TRACKING_URI, BaseAction) +# from gaiaflow.managers.base_manager import BaseGaiaflowManager +# +# +# class MlflowModelManager(BaseGaiaflowManager): +# def __init__( +# self, +# registry_uri=None, +# action: BaseAction = None, +# force_new: bool = False, +# prune: bool = False, +# local: bool = False, +# **kwargs, +# ): +# self.docker_client = docker.from_env() +# self.registry_uri = registry_uri.rstrip("/") if registry_uri else None +# self.local = local +# +# # TODO: not optimal. Use ContextManager to restore the env vars +# # Move them to the actual methods where they are needed. +# os.environ["MLFLOW_TRACKING_URI"] = ( +# os.getenv("MLFLOW_TRACKING_URI") or "https://localhost:5000" +# ) +# os.environ["MLFLOW_S3_ENDPOINT_URL"] = ( +# os.getenv("MLFLOW_S3_ENDPOINT_URL") or "https://localhost:9000" +# ) +# os.environ["AWS_ACCESS_KEY_ID"] = ( +# os.getenv("AWS_ACCESS_KEY_ID") or DEFAULT_MINIO_AWS_ACCESS_KEY_ID +# ) +# os.environ["AWS_SECRET_ACCESS_KEY"] = ( +# os.getenv("AWS_SECRET_ACCESS_KEY") or DEFAULT_MINIO_AWS_SECRET_ACCESS_KEY +# ) +# +# mlflow_tracking_uri = ( +# os.getenv("MLFLOW_TRACKING_URI") or DEFAULT_MLFLOW_TRACKING_URI +# ) +# mlflow.set_tracking_uri(mlflow_tracking_uri) +# +# # This Manager does not need the paths. But it expects that either +# # the model exists in S3 if s3_uri is used or a live instance of MLFlow +# # running if run_id is used +# super().__init__( +# gaiaflow_path="", +# user_project_path="", +# action=action, +# force_new=force_new, +# prune=prune, +# **kwargs, +# ) +# +# def _get_model_uri(self, model_uri=None, run_id=None): +# if model_uri: +# return model_uri +# elif run_id: +# return f"runs:/{run_id}/model" +# else: +# raise ValueError("Either model_uri or run_id must be provided.") +# +# # TODO: Use this instead +# # @classmethod +# # def run(cls, action, **kwargs): +# # params = kwargs.get("params", {}) +# # model_uri = params.get("model_uri") +# # run_id = params.get("run_id") +# # image_name = params.get("image_name") +# # enable_mlserver = params.get("enable_mlserver", True) +# # manag = MlflowModelManager(model_uri, run_id, image_name, +# # enable_mlserver) +# # environment = create_env() +# # with set_env_cm(environment): +# # if action == BaseAction.START: +# # manag.start() +# +# def start(self, **kwargs): +# print("kwargs inside start, **kwargs", kwargs) +# params = kwargs.get("params", {}) +# +# model_uri = params.get("model_uri") +# run_id = params.get("run_id") +# image_name = params.get("image_name") +# enable_mlserver = params.get("enable_mlserver", True) +# +# self.build_model_docker_image(model_uri, run_id, image_name, enable_mlserver) +# self.run_container(image_name) +# +# def stop(self, **kwargs): +# print("kwargs inside stop, **kwargs", kwargs) +# params = kwargs.get("params", {}) +# +# container_id_or_name = params.get("container_id_or_name") +# self.stop_and_remove_container(container_id_or_name) +# +# def cleanup(self, **kwargs): +# print("kwargs inside cleanup, **kwargs", kwargs) +# params = kwargs.get("params", {}) +# +# container_id_or_name = params.get("container_id_or_name") +# purge = params.get("purge") +# self.stop_and_remove_container(container_id_or_name, purge) +# +# def build_model_docker_image( +# self, model_uri=None, run_id=None, image_name=None, enable_mlserver=True +# ): +# model_uri_final = self._get_model_uri(model_uri, run_id) +# +# unique_tag = str(uuid.uuid4())[:8] +# image_name = image_name or f"mlflow-model:{unique_tag}" +# +# full_image_name = ( +# f"{self.registry_uri}/{image_name}" if self.registry_uri else image_name +# ) +# +# print(f"Building Docker image from {model_uri_final} as {full_image_name}...") +# +# mlflow.models.build_docker( +# model_uri=model_uri_final, +# name=full_image_name, +# enable_mlserver=enable_mlserver, +# ) +# +# return full_image_name +# +# def run_container(self, image_name, port=8246): +# container = self.docker_client.containers.run( +# image_name, +# detach=True, +# ports={"8080/tcp": port}, +# name=f"mlflow_model_{uuid.uuid4().hex[:6]}", +# ) +# print(f"Container started: {container.name} on port {port}") +# return container +# +# def stop_and_remove_container(self, container_id_or_name, purge=False): +# try: +# container = self.docker_client.containers.get(container_id_or_name) +# print(f"Stopping container: {container.name}") +# container.stop() +# container.remove() +# print("Container stopped and removed.") +# if purge: +# image_name = container.image.tags[0] +# self.docker_client.images.remove(image=image_name, force=True) +# print(f"Image {image_name} removed.") +# except docker.context.api.errors.NotFound: +# print("Container not found.") +# +# def push_image_to_registry(self, image_name): +# if not self.registry_uri: +# raise ValueError("No registry_uri provided.") +# print(f"Pushing image {image_name} to registry {self.registry_uri}...") +# subprocess.run(["docker", "push", image_name], check=True) +# print("Image pushed.") +# +# def generate_k8s_yaml(self): +# print(" Generating Kubernetes deployment YAML...") +# raise NotImplementedError() +# +# def deploy_to_kubernetes(self, k8s_yaml_path): +# print(f"Deploying to Kubernetes from {k8s_yaml_path}...") +# subprocess.run(["kubectl", "apply", "-f", k8s_yaml_path], check=True) +# print("Deployment applied.") +# +# +# # +# # @contextmanager +# # def set_env_cm(**new_env: str | None) -> Generator[dict[str, str | None]]: +# # """Run the code in the block with a new environment `new_env`.""" +# # restore_env = set_env(**new_env) +# # try: +# # yield new_env +# # finally: +# # restore_env() +# # +# # +# # def set_env(**new_env: str | None) -> Callable[[], None]: +# # """ +# # Set the new environment in `new_env` and return a no-arg +# # function to restore the old environment. +# # """ +# # old_env = {k: os.environ.get(k) for k in new_env.keys()} +# # +# # def restore_env(): +# # for ko, vo in old_env.items(): +# # if vo is not None: +# # os.environ[ko] = vo +# # elif ko in os.environ: +# # del os.environ[ko] +# # +# # for kn, vn in new_env.items(): +# # if vn is not None: +# # os.environ[kn] = vn +# # elif kn in os.environ: +# # del os.environ[kn] +# # +# # return restore_env diff --git a/src/gaiaflow/managers/mlops_manager.py b/src/gaiaflow/managers/mlops_manager.py index 7c87ef1..6929c22 100644 --- a/src/gaiaflow/managers/mlops_manager.py +++ b/src/gaiaflow/managers/mlops_manager.py @@ -19,41 +19,172 @@ MLFLOW_SERVICES, Action, BaseAction, - Service, ExtendedAction, + Service, ) from gaiaflow.managers.base_manager import BaseGaiaflowManager -from gaiaflow.managers.utils import (create_directory, delete_project_state, - find_python_packages, - gaiaflow_path_exists_in_state, - handle_error, log_error, log_info, run, - save_project_state, set_permissions, convert_crlf_to_lf, - env_exists, - update_micromamba_env_in_docker) - -_IMAGES = [ - "docker-compose-airflow-apiserver:latest", - "docker-compose-airflow-scheduler:latest", - "docker-compose-airflow-dag-processor:latest", - "docker-compose-airflow-triggerer:latest", - "docker-compose-airflow-init:latest", - "docker-compose-mlflow:latest", - "minio/mc:latest", - "minio/minio:latest", - "postgres:13", -] - -_AIRFLOW_CONTAINERS = [ - "airflow-apiserver", - "airflow-scheduler", - "airflow-dag-processor", - "airflow-triggerer" -] - -_VOLUMES = [ - "docker-compose_postgres-db-volume-airflow", - "docker-compose_postgres-db-volume-mlflow", -] +from gaiaflow.managers.utils import ( + convert_crlf_to_lf, + create_directory, + delete_project_state, + env_exists, + find_python_packages, + gaiaflow_path_exists_in_state, + handle_error, + log_error, + log_info, + run, + save_project_state, + set_permissions, + update_micromamba_env_in_docker, +) + + +class DockerResources: + IMAGES = [ + "docker-compose-airflow-apiserver:latest", + "docker-compose-airflow-scheduler:latest", + "docker-compose-airflow-dag-processor:latest", + "docker-compose-airflow-triggerer:latest", + "docker-compose-airflow-init:latest", + "docker-compose-mlflow:latest", + "minio/mc:latest", + "minio/minio:latest", + "postgres:13", + ] + + AIRFLOW_CONTAINERS = [ + "airflow-apiserver", + "airflow-scheduler", + "airflow-dag-processor", + "airflow-triggerer", + ] + + VOLUMES = [ + "docker-compose_postgres-db-volume-airflow", + "docker-compose_postgres-db-volume-mlflow", + ] + + SERVICES = { + "airflow": AIRFLOW_SERVICES, + "mlflow": MLFLOW_SERVICES, + "minio": MINIO_SERVICES, + } + + +class DockerHelper: + def __init__(self, gaiaflow_path: Path, is_prod_local: bool): + self.gaiaflow_path = gaiaflow_path + self.is_prod_local = is_prod_local + + def _base_cmd(self) -> list[str]: + base = [ + "docker", + "compose", + "-f", + f"{self.gaiaflow_path}/docker_stuff/docker-compose/docker-compose.yml", + ] + if self.is_prod_local: + base += [ + "-f", + f"{self.gaiaflow_path}/docker_stuff/docker-compose/docker-compose-minikube-network.yml", + ] + return base + + @staticmethod + def docker_services_for(component: str) -> list[str]: + return DockerResources.SERVICES.get(component, []) + + def run_compose(self, actions: list[str], service: str | None = None): + cmd = self._base_cmd() + if service: + services = self.docker_services_for(service) + if not services: + handle_error(f"Unknown service: {service}") + cmd += actions + services + else: + cmd += actions + + log_info(f"Running: {' '.join(cmd)}") + run(cmd, f"Error running docker compose {actions}") + + @staticmethod + def prune(): + prune_cmds = [ + ( + ["docker", "builder", "prune", "-a", "-f"], + "Error pruning docker build cache", + ), + (["docker", "system", "prune", "-a", "-f"], "Error pruning docker system"), + (["docker", "volume", "prune", "-a", "-f"], "Error pruning docker volumes"), + ( + ["docker", "network", "rm", "docker-compose_ml-network"], + "Error removing docker network", + ), + ] + for cmd, msg in prune_cmds: + run(cmd, msg) + + for image in DockerResources.IMAGES: + run(["docker", "rmi", "-f", image], f"Error deleting image {image}") + for volume in DockerResources.VOLUMES: + run(["docker", "volume", "rm", volume], f"Error removing volume {volume}") + + +class JupyterHelper: + def __init__( + self, port: int, env_tool: str, user_env_name: str | None, gaiaflow_path: Path + ): + self.port = port + self.env_tool = env_tool + self.user_env_name = user_env_name + self.gaiaflow_path = gaiaflow_path + + def check_port(self): + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + if sock.connect_ex(("127.0.0.1", self.port)) == 0: + handle_error(f"Port {self.port} is already in use.") + + def stop(self): + log_info(f"Attempting to stop Jupyter processes on port {self.port}") + for proc in psutil.process_iter(attrs=["pid", "name", "cmdline"]): + try: + cmdline = proc.info.get("cmdline") or [] + name = proc.info.get("name") or "" + if "jupyter" in name or any("jupyter-lab" in arg for arg in cmdline): + log_info(f"Terminating process {proc.pid} ({name})") + proc.terminate() + proc.wait(timeout=5) + except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess): + continue + + def start(self): + env_name = self.get_env_name() + if not env_exists(env_name, env_tool=self.env_tool): + print( + f"Environment {env_name} not found. Run `mamba env create -f environment.yml`?" + ) + return + cmd = [ + self.env_tool, + "run", + "-n", + env_name, + "jupyter", + "lab", + "--ip=0.0.0.0", + f"--port={self.port}", + ] + log_info("Starting Jupyter Lab..." + " ".join(cmd)) + subprocess.Popen(cmd) + + def get_env_name(self): + if self.user_env_name: + return self.user_env_name + env_path = Path(self.gaiaflow_path).resolve() / "environment.yml" + with open(env_path, "r") as f: + env_yml = yaml.safe_load(f) + return env_yml.get("name") class MlopsManager(BaseGaiaflowManager): @@ -84,16 +215,20 @@ def __init__( ) self.service = service self.cache = cache - self.jupyter_port = jupyter_port self.delete_volume = delete_volume self.docker_build = docker_build self.os_type = platform.system().lower() - self.project_root = Path(__file__).resolve().parent + # self.project_root = Path(__file__).resolve().parent self.fs = fsspec.filesystem("file") self.prod_local = prod_local self.user_env_name = user_env_name self.env_tool = env_tool + self.docker = DockerHelper(gaiaflow_path, prod_local) + self.jupyter = JupyterHelper( + jupyter_port, env_tool, user_env_name, gaiaflow_path + ) + super().__init__( gaiaflow_path=gaiaflow_path, user_project_path=user_project_path, @@ -108,7 +243,7 @@ def run(cls, **kwargs): if action is None: raise ValueError("Missing required argument 'action'") - manager = MlopsManager(**kwargs) + manager = cls(**kwargs) action_map = { BaseAction.START: manager.start, @@ -119,35 +254,18 @@ def run(cls, **kwargs): } try: - action_method = action_map[action] + action_map[action]() except KeyError: raise ValueError(f"Unknown action: {action}") - action_method() - def start(self): log_info("Starting Gaiaflow services") if self.force_new: self.cleanup() - gaiaflow_path_exists = gaiaflow_path_exists_in_state(self.gaiaflow_path, True) - - if not gaiaflow_path_exists: - log_info("Setting up directories...") - create_directory(f"{self.user_project_path}/logs") - create_directory(f"{self.user_project_path}/data") - - log_info("Updating .env file...") - self._update_env_file_with_airflow_uid(f"{self.user_project_path}/.env") - - log_info("Creating gaiaflow context...") - self._create_gaiaflow_context() - - log_info("Updating gaiaflow context with user project information...") - self._update_files() - - save_project_state(self.user_project_path, self.gaiaflow_path) + if not gaiaflow_path_exists_in_state(self.gaiaflow_path, True): + self._setup_project_context() else: log_info( "Gaiaflow project already exists at " @@ -157,144 +275,106 @@ def start(self): self._copy_user_env_file() - if self.service == Service.jupyter or self.service == Service.all: - self._check_port() + if self.service in {Service.jupyter, Service.all}: + self.jupyter.check_port() if self.docker_build: - build_cmd = ["build"] - if not self.cache: - build_cmd.append("--no-cache") - - log_info("Building Docker images") - - if self.service == Service.all: - self._docker_compose_action(build_cmd, service=None) - elif self.service == Service.jupyter: - pass - else: - self._docker_compose_action(build_cmd, self.service) + self._build_docker_images() - if self.service == Service.all: - self._start_jupyter() - self._docker_compose_action(["up", "-d"], service=None) - elif self.service == Service.jupyter: - self._start_jupyter() - else: - self._docker_compose_action(["up", "-d"], service=self.service) + self._start_services() def stop(self): log_info("Shutting down Gaiaflow services...") + if self.service == Service.jupyter: - self._stop_jupyter() + self.jupyter.stop() elif self.service == Service.all: - down_cmd = ["down"] - if self.delete_volume: - log_info("Removing volumes with shutdown") - down_cmd.append("-v") - self._docker_compose_action(down_cmd) - self._stop_jupyter() + self._stop_all_services() else: - down_cmd = ["down"] - if self.delete_volume: - log_info("Removing volumes with shutdown") - down_cmd.append("-v") - self._docker_compose_action(down_cmd, self.service) + self._stop_service(self.service) log_info("Stopped Gaiaflow services successfully") + def cleanup(self): + try: + log_info(f"Attempting deleting Gaiaflow context at {self.gaiaflow_path}") + shutil.rmtree(self.gaiaflow_path) + except FileNotFoundError: + log_error(f"Gaiaflow context not found at {self.gaiaflow_path}") + + try: + log_info( + f"Attempting deleting Gaiaflow project state at {GAIAFLOW_STATE_FILE}" + ) + delete_project_state(self.gaiaflow_path) + except (json.JSONDecodeError, FileNotFoundError): + raise + if self.prune: + self.docker.prune() + + log_info("Gaiaflow cleanup complete!") + @staticmethod def update_deps(): log_info("Running update_deps") - update_micromamba_env_in_docker(_AIRFLOW_CONTAINERS) + update_micromamba_env_in_docker(DockerResources.AIRFLOW_CONTAINERS) log_info("Finished running update_deps") - def _check_port(self): - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: - if sock.connect_ex(("127.0.0.1", self.jupyter_port)) == 0: - handle_error(f"Port {self.jupyter_port} is already in use.") + def _setup_project_context(self): + create_directory(self.user_project_path / "logs") + create_directory(self.user_project_path / "data") - def _stop_jupyter(self): - log_info(f"Attempting to stop Jupyter processes on port {self.jupyter_port}") - for proc in psutil.process_iter(attrs=["pid", "name", "cmdline"]): - try: - cmdline = proc.info.get("cmdline") or [] - name = proc.info.get("name") or "" - if any("jupyter-lab" in arg for arg in cmdline) or "jupyter" in name: - log_info(f"Terminating process {proc.pid} ({name})") - proc.terminate() - proc.wait(timeout=5) - except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess): - continue - - @staticmethod - def _docker_services_for(component): - services = { - "airflow": AIRFLOW_SERVICES, - "mlflow": MLFLOW_SERVICES, - "minio": MINIO_SERVICES, - } - return services.get(component, []) - - def _docker_compose_action(self, actions, service=None): - if self.prod_local: - base_cmd = [ - "docker", - "compose", - "-f", - f"{self.gaiaflow_path}/docker_stuff/docker-compose/docker-compose.yml", - "-f", - f"{self.gaiaflow_path}/docker_stuff/docker-compose/docker-compose-minikube-network.yml", - ] - else: - base_cmd = [ - "docker", - "compose", - "-f", - f"{self.gaiaflow_path}/docker_stuff/docker-compose/docker-compose.yml", - ] - if service: - services = MlopsManager._docker_services_for(service) - print("Services:::", services, service) - if not services: - handle_error(f"Unknown service: {service}") - cmd = base_cmd + actions + services - else: - cmd = base_cmd + actions + log_info("Updating .env file...") + self._update_env_file_with_airflow_uid(self.user_project_path / ".env") - log_info(f"Running: {' '.join(cmd)}") - run(cmd, f"Error running docker compose {actions}") + log_info("Creating gaiaflow context...") + self._create_gaiaflow_context() - def get_env_name(self): - if self.user_env_name: - return self.user_env_name + log_info("Updating gaiaflow context with user project information...") + self._update_files() - ctx_path = Path(self.gaiaflow_path).resolve() - env_path = ctx_path / "environment.yml" - with open(env_path, "r") as f: - env_yml = yaml.safe_load(f) - return env_yml.get("name") + save_project_state(self.user_project_path, self.gaiaflow_path) - def _start_jupyter(self): - env_name = self.get_env_name() - if not env_exists(env_name, env_tool=self.env_tool): - print( - f"Environment {env_name} not found. Run `mamba env create " - f"-f environment.yml`?" - ) - cmd = [self.env_tool, "run", "-n", env_name, "jupyter", "lab", "--ip=0.0.0.0", f"--port={self.jupyter_port}"] - log_info("Starting Jupyter Lab..." + " ".join(cmd)) - subprocess.Popen(cmd) + def _build_docker_images(self): + build_cmd = ["build"] + if not self.cache: + build_cmd.append("--no-cache") + log_info("Building Docker images") + if self.service == Service.all: + self.docker.run_compose(build_cmd) + elif self.service != Service.jupyter: + self.docker.run_compose(build_cmd, self.service) - def _update_env_file_with_airflow_uid(self, env_path): - if self.os_type == "linux": - uid = str(os.getuid()) + def _start_services(self): + if self.service == Service.all: + self.jupyter.start() + self.docker.run_compose(["up", "-d"]) + elif self.service == Service.jupyter: + self.jupyter.start() else: - uid = 50000 + self.docker.run_compose(["up", "-d"], self.service) + + def _stop_all_services(self): + down_cmd = ["down"] + if self.delete_volume: + log_info("Removing volumes with shutdown") + down_cmd.append("-v") + self.docker.run_compose(down_cmd) + self.jupyter.stop() + + def _stop_service(self, service: Service): + down_cmd = ["down"] + if self.delete_volume: + log_info("Removing volumes with shutdown") + down_cmd.append("-v") + self.docker.run_compose(down_cmd, service) + + def _update_env_file_with_airflow_uid(self, env_path: Path): + uid = str(os.getuid()) if self.os_type == "linux" else "50000" lines = [] - if os.path.exists(env_path): - with open(env_path, "r") as f: - lines = f.readlines() + if env_path.exists(): + lines = env_path.read_text().splitlines(keepends=True) key_found = False new_lines = [] @@ -308,8 +388,7 @@ def _update_env_file_with_airflow_uid(self, env_path): if not key_found: new_lines.append(f"AIRFLOW_UID={uid}\n") - with open(env_path, "w") as f: - f.writelines(new_lines) + env_path.write_text("".join(new_lines)) log_info(f"Set AIRFLOW_UID={uid} in {env_path}") @@ -325,69 +404,83 @@ def _create_gaiaflow_context(self): package_dir = Path(__file__).parent.parent.resolve() docker_dir = package_dir.parent / "docker_stuff" - # docker_dir = package_dir / "docker_stuff" - shutil.copytree(docker_dir, self.gaiaflow_path / "docker_stuff", dirs_exist_ok=True) - self._copy_user_env_file() - log_info(f"Gaiaflow context created at {self.gaiaflow_path}") - - def _update_files(self): - yaml = YAML() - yaml.preserve_quotes = True - - compose_path = ( - self.gaiaflow_path / "docker_stuff" / "docker-compose" / "docker-compose.yml" + shutil.copytree( + docker_dir, self.gaiaflow_path / "docker_stuff", dirs_exist_ok=True ) + log_info(f"Gaiaflow context created at {self.gaiaflow_path}") - with open(compose_path) as f: - compose_data = yaml.load(f) - + def _collect_volumes(self, compose_data: dict) -> list[str]: x_common = compose_data.get("x-airflow-common", {}) original_vols = x_common.get("volumes", []) new_volumes = [] + + # Re-map predefined volumes to absolute paths for vol in original_vols: - parts = vol.split(":", 1) - if len(parts) == 2: - src, dst = parts - src_path = (self.user_project_path / Path(src).name).resolve().as_posix() - print("src_path", src_path) + if ":" in vol: + src, dst = vol.split(":", 1) + src_path = ( + (self.user_project_path / Path(src).name).resolve().as_posix() + ) new_volumes.append(f"{src_path}:{dst}") + # Collect User Python packages from their project directory to mount + # them to docker containers existing_mounts = {Path(v.split(":", 1)[0]).name for v in new_volumes} python_packages = find_python_packages(self.user_project_path) - for python_package in python_packages: - set_permissions(python_package, 0o755) + # Set permissions so that docker containers can execute the code in + # their package + for package in python_packages: + set_permissions(package, 0o755) for child in self.user_project_path.iterdir(): if ( - child.is_dir() and child.name not in existing_mounts and child.name - ) in python_packages: + child.is_dir() + and child.name not in existing_mounts + and child.name in python_packages + ): dst_path = f"/opt/airflow/{child.name}" new_volumes.append(f"{child.resolve().as_posix()}:{dst_path}") - kube_config_path = (self.gaiaflow_path.resolve() / "docker_stuff" / "kube_config_inline").as_posix() - new_volumes.append(f"{kube_config_path}:/home/airflow/.kube/config") - - entrypoint_path = ( - self.gaiaflow_path.resolve() / "docker_stuff" / "docker-compose" / "entrypoint.sh").as_posix() - new_volumes.append(f"{entrypoint_path}:/opt/airflow/entrypoint.sh") + # Add special mounts for prod_local mode + kube_config = ( + self.gaiaflow_path.resolve() / "docker_stuff" / "kube_config_inline" + ).as_posix() + entrypoint = ( + self.gaiaflow_path.resolve() + / "docker_stuff" + / "docker-compose" + / "entrypoint.sh" + ).as_posix() + pyproject = (self.user_project_path.resolve() / "pyproject.toml").as_posix() + env_file = (self.user_project_path.resolve() / "environment.yml").as_posix() - # new_volumes.append( - # f"{self.gaiaflow_path.as_posix() / 'docker_stuff' / 'docker-compose'}/entrypoint.sh:/opt/airflow/entrypoint.sh" - # ) + new_volumes += [ + f"{kube_config}:/home/airflow/.kube/config", + f"{entrypoint}:/opt/airflow/entrypoint.sh", + f"{pyproject}:/opt/airflow/pyproject.toml", + f"{env_file}:/opt/airflow/environment.yml", + "/var/run/docker.sock:/var/run/docker.sock", + ] - pyproject_path = (self.user_project_path.resolve() / "pyproject.toml").as_posix() - new_volumes.append(f"{pyproject_path}:/opt/airflow/pyproject.toml") + return new_volumes - pyproject_path = ( - self.user_project_path.resolve() / "environment.yml" - ).as_posix() + def _update_files(self): + yaml = YAML() + yaml.preserve_quotes = True - new_volumes.append(f"{pyproject_path}:/opt/airflow/environment.yml") + compose_path = ( + self.gaiaflow_path + / "docker_stuff" + / "docker-compose" + / "docker-compose.yml" + ) - new_volumes.append("/var/run/docker.sock:/var/run/docker.sock") + with open(compose_path) as f: + compose_data = yaml.load(f) + new_volumes = self._collect_volumes(compose_data) compose_data["x-airflow-common"]["volumes"] = new_volumes with compose_path.open("w") as f: @@ -399,47 +492,5 @@ def _update_files(self): set_permissions(entrypoint_path) convert_crlf_to_lf(entrypoint_path) - def cleanup(self): - try: - log_info(f"Attempting deleting Gaiaflow context at { - self.gaiaflow_path}") - shutil.rmtree(self.gaiaflow_path) - except FileNotFoundError: - log_error(f"Gaiaflow context not found at {self.gaiaflow_path}") - try: - log_info( - f"Attempting deleting Gaiaflow project state at {GAIAFLOW_STATE_FILE}" - ) - delete_project_state(self.gaiaflow_path) - except (json.JSONDecodeError, FileNotFoundError): - raise - if self.prune: - run( - ["docker", "builder", "prune", "-a", "-f"], - "Error pruning docker build cache", - ) - run( - ["docker", "system", "prune", "-a", "-f"], - "Error pruning docker system", - ) - run( - ["docker", "volume", "prune", "-a", "-f"], - "Error pruning docker volumes", - ) - run( - ["docker", "network", "rm", "docker-compose_ml-network"], - "Error removing docker network", - ) - for image in _IMAGES: - run(["docker", "rmi", "-f", image], f"Error deleting image {image}") - for volume in _VOLUMES: - run( - ["docker", "volume", "rm", volume], - f"Error removing volume {volume}", - ) - log_info("Gaiaflow cleanup complete!") - def _get_valid_actions(self) -> Set[Action]: - base_actions = super()._get_valid_actions() - extra_actions = {ExtendedAction.UPDATE_DEPS} - return base_actions | extra_actions + return super()._get_valid_actions() | {ExtendedAction.UPDATE_DEPS} diff --git a/src/gaiaflow/managers/utils.py b/src/gaiaflow/managers/utils.py index 429ff3e..53a0f52 100644 --- a/src/gaiaflow/managers/utils.py +++ b/src/gaiaflow/managers/utils.py @@ -1,4 +1,3 @@ -import docker import json import subprocess import sys @@ -8,6 +7,7 @@ from datetime import datetime from pathlib import Path +import docker import fsspec import typer @@ -16,7 +16,7 @@ fs = fsspec.filesystem("file") -def get_gaialfow_version() -> str: +def get_gaiaflow_version() -> str: try: from importlib.metadata import version @@ -44,7 +44,10 @@ def log_info(message: str): def log_error(message: str): - print(f"\033[0;31mERROR:\033[0m {message}", file=sys.stderr) + print( + f"\033[0;31m[{datetime.now().strftime('%H:%M:%S')}]ERROR:\033[0m {message}", + file=sys.stderr, + ) def run(command: list, error_message: str, env=None): @@ -173,7 +176,7 @@ def gaiaflow_path_exists_in_state(gaiaflow_path: Path, check_fs: bool = True) -> def delete_project_state(gaiaflow_path: Path): state_file = get_state_file() - print("state_file", state_file) + log_info("state_file: " + str(state_file)) if not state_file.exists(): log_error( "State file not found at ~/.gaiaflow/state.json. Please run the services." @@ -184,13 +187,13 @@ def delete_project_state(gaiaflow_path: Path): with open(state_file, "r") as f: state = json.load(f) - print("found!", state.get("gaiaflow_path"), state) + log_info("found! " + str(state.get("gaiaflow_path")) + str(state)) key = str(gaiaflow_path) if key in state: del state[key] with open(state_file, "w") as f: json.dump(state, f, indent=2) - except (json.JSONDecodeError, FileNotFoundError, Exception): + except (json.JSONDecodeError, FileNotFoundError, AttributeError, Exception): raise @@ -229,22 +232,16 @@ def create_gaiaflow_context_path(project_path: Path) -> tuple[Path, Path]: user_project_path = Path(project_path).resolve() if not user_project_path.exists(): raise FileNotFoundError(f"{user_project_path} not found") - version = get_gaialfow_version() + version = get_gaiaflow_version() # project_name = str(user_project_path).split("/")[-1] project_name = user_project_path.name tmp_dir = Path(tempfile.gettempdir()) - gaiaflow_path = tmp_dir / f"gaiaflow-{version}-{project_name}" + gaiaflow_path = tmp_dir / f"gaiaflow-{version}-{project_name}" return gaiaflow_path, user_project_path def convert_crlf_to_lf(file_path: str): - """ - Converts a file from Windows-style CRLF line endings to Unix-style LF line endings. - - Args: - file_path (str): Path to the file to convert. - """ with open(file_path, "rb") as f: content = f.read() @@ -264,6 +261,7 @@ def is_wsl() -> bool: except FileNotFoundError: return False + def env_exists(env_name, env_tool="mamba"): result = subprocess.run( [env_tool, "env", "list", "--json"], capture_output=True, text=True @@ -271,11 +269,12 @@ def env_exists(env_name, env_tool="mamba"): envs = json.loads(result.stdout).get("envs", []) return any(env_name in env for env in envs) + def update_micromamba_env_in_docker( - containers: list[str], - env_name: str = "default_user_env", - max_workers: int = 4, - ): + containers: list[str], + env_name: str = "default_user_env", + max_workers: int = 4, +): client = docker.from_env() def _update_one(cname: str): @@ -297,13 +296,11 @@ def _update_one(cname: str): log_info(output.decode()) with ThreadPoolExecutor(max_workers=max_workers) as executor: - futures = { - executor.submit(_update_one, cname): cname for cname in containers - } + futures = {executor.submit(_update_one, cname): cname for cname in containers} for future in as_completed(futures): cname = futures[future] try: future.result() except Exception as e: - log_error(f"[{cname}] Unexpected error: {e}") \ No newline at end of file + log_error(f"[{cname}] Unexpected error: {e}") diff --git a/src/gaiaflow/testing.py b/src/gaiaflow/testing.py new file mode 100644 index 0000000..b5a0e2f --- /dev/null +++ b/src/gaiaflow/testing.py @@ -0,0 +1,37 @@ +import os +from collections.abc import Generator +from contextlib import contextmanager +from typing import Callable + + +@contextmanager +def set_env_cm(**new_env: str | None) -> Generator[dict[str, str | None]]: + """Run the code in the block with a new environment `new_env`.""" + restore_env = set_env(**new_env) + try: + yield new_env + finally: + restore_env() + + +def set_env(**new_env: str | None) -> Callable[[], None]: + """ + Set the new environment in `new_env` and return a no-arg + function to restore the old environment. + """ + old_env = {k: os.environ.get(k) for k in new_env.keys()} + + def restore_env(): + for ko, vo in old_env.items(): + if vo is not None: + os.environ[ko] = vo + elif ko in os.environ: + del os.environ[ko] + + for kn, vn in new_env.items(): + if vn is not None: + os.environ[kn] = vn + elif kn in os.environ: + del os.environ[kn] + + return restore_env diff --git a/tests/cli/__init__.py b/tests/cli/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/cli/commands/__init__.py b/tests/cli/commands/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/cli/commands/test_minikube.py b/tests/cli/commands/test_minikube.py new file mode 100644 index 0000000..269bf76 --- /dev/null +++ b/tests/cli/commands/test_minikube.py @@ -0,0 +1,256 @@ +import unittest +from unittest.mock import patch, MagicMock +from pathlib import Path +from types import SimpleNamespace +from typer.testing import CliRunner + +from gaiaflow.constants import DEFAULT_IMAGE_NAME, BaseAction, \ + ExtendedAction +from gaiaflow.cli.commands.minikube import app as prod_app +from gaiaflow.cli.commands.minikube import load_imports + +class TestGaiaflowProdCLI(unittest.TestCase): + def setUp(self): + self.runner = CliRunner() + self.test_project_path = Path("/test/project") + self.test_gaiaflow_path = Path("/tmp/gaiaflow/test") + + self.mock_imports = SimpleNamespace( + BaseAction=BaseAction, + ExtendedAction=ExtendedAction, + MinikubeManager=MagicMock(), + create_gaiaflow_context_path=MagicMock( + return_value=(self.test_gaiaflow_path, self.test_project_path) + ), + gaiaflow_path_exists_in_state=MagicMock(return_value=True), + parse_key_value_pairs=MagicMock(return_value={"key1": "value1", "key2": "value2"}), + ) + + @patch('gaiaflow.cli.commands.minikube.load_imports') + def test_start_command(self, mock_load_imports): + mock_load_imports.return_value = self.mock_imports + + result = self.runner.invoke(prod_app, [ + "start", + "--path", str(self.test_project_path), + "--force-new" + ]) + + self.assertEqual(result.exit_code, 0) + + self.mock_imports.create_gaiaflow_context_path.assert_called_once_with( + self.test_project_path + ) + self.mock_imports.gaiaflow_path_exists_in_state.assert_called_once_with( + self.test_gaiaflow_path, True + ) + + self.mock_imports.MinikubeManager.run.assert_called_once_with( + gaiaflow_path=self.test_gaiaflow_path, + user_project_path=self.test_project_path, + action=BaseAction.START, + force_new=True, + ) + + @patch('gaiaflow.cli.commands.minikube.load_imports') + def test_start_command_exits_when_project_not_exists(self, mock_load_imports): + mock_load_imports.return_value = self.mock_imports + self.mock_imports.gaiaflow_path_exists_in_state.return_value = False + + result = self.runner.invoke(prod_app, [ + "start", + "--path", str(self.test_project_path) + ]) + + self.assertEqual(result.exit_code, 0) + self.assertIn("Please create a project with Gaiaflow", result.output) + + self.mock_imports.MinikubeManager.run.assert_not_called() + + @patch('gaiaflow.cli.commands.minikube.load_imports') + def test_stop_command(self, mock_load_imports): + mock_load_imports.return_value = self.mock_imports + + result = self.runner.invoke(prod_app, [ + "stop", + "--path", str(self.test_project_path) + ]) + + self.assertEqual(result.exit_code, 0) + + self.mock_imports.MinikubeManager.run.assert_called_once_with( + gaiaflow_path=self.test_gaiaflow_path, + user_project_path=self.test_project_path, + action=BaseAction.STOP, + ) + + @patch('gaiaflow.cli.commands.minikube.load_imports') + def test_restart_command(self, mock_load_imports): + mock_load_imports.return_value = self.mock_imports + + result = self.runner.invoke(prod_app, [ + "restart", + "--path", str(self.test_project_path), + "--force-new" + ]) + + self.assertEqual(result.exit_code, 0) + + self.mock_imports.MinikubeManager.run.assert_called_once_with( + gaiaflow_path=self.test_gaiaflow_path, + user_project_path=self.test_project_path, + action=BaseAction.RESTART, + ) + + @patch('gaiaflow.cli.commands.minikube.load_imports') + def test_dockerize_command(self, mock_load_imports): + mock_load_imports.return_value = self.mock_imports + + result = self.runner.invoke(prod_app, [ + "dockerize", + "--path", str(self.test_project_path), + "--image-name", "my-custom-image" + ]) + + self.assertEqual(result.exit_code, 0) + + self.mock_imports.MinikubeManager.run.assert_called_once_with( + gaiaflow_path=self.test_gaiaflow_path, + user_project_path=self.test_project_path, + action=ExtendedAction.DOCKERIZE, + local=False, + image_name="my-custom-image" + ) + + @patch('gaiaflow.cli.commands.minikube.load_imports') + def test_dockerize_command_with_default_image_name(self, mock_load_imports): + mock_load_imports.return_value = self.mock_imports + + result = self.runner.invoke(prod_app, [ + "dockerize", + "--path", str(self.test_project_path) + ]) + + self.assertEqual(result.exit_code, 0) + + self.mock_imports.MinikubeManager.run.assert_called_once() + call_args = self.mock_imports.MinikubeManager.run.call_args + self.assertEqual(call_args.kwargs['image_name'], DEFAULT_IMAGE_NAME) + + @patch('gaiaflow.cli.commands.minikube.load_imports') + def test_create_config_command(self, mock_load_imports): + mock_load_imports.return_value = self.mock_imports + + result = self.runner.invoke(prod_app, [ + "create-config", + "--path", str(self.test_project_path) + ]) + + self.assertEqual(result.exit_code, 0) + + self.mock_imports.MinikubeManager.run.assert_called_once_with( + gaiaflow_path=self.test_gaiaflow_path, + user_project_path=self.test_project_path, + action=ExtendedAction.CREATE_CONFIG, + ) + + @patch('gaiaflow.cli.commands.minikube.load_imports') + def test_create_secret_command(self, mock_load_imports): + mock_load_imports.return_value = self.mock_imports + + result = self.runner.invoke(prod_app, [ + "create-secret", + "--path", str(self.test_project_path), + "--name", "my-secret", + "--data", "key1=value1", + "--data", "key2=value2" + ]) + + self.assertEqual(result.exit_code, 0) + + self.mock_imports.parse_key_value_pairs.assert_called_once_with( + ["key1=value1", "key2=value2"] + ) + + self.mock_imports.MinikubeManager.run.assert_called_once_with( + gaiaflow_path=self.test_gaiaflow_path, + user_project_path=self.test_project_path, + action=ExtendedAction.CREATE_SECRET, + secret_name="my-secret", + secret_data={"key1": "value1", "key2": "value2"}, + ) + + @patch('gaiaflow.cli.commands.minikube.load_imports') + def test_cleanup_command(self, mock_load_imports): + mock_load_imports.return_value = self.mock_imports + + result = self.runner.invoke(prod_app, [ + "cleanup", + "--path", str(self.test_project_path) + ]) + + self.assertEqual(result.exit_code, 0) + + self.mock_imports.MinikubeManager.run.assert_called_once_with( + gaiaflow_path=self.test_gaiaflow_path, + user_project_path=self.test_project_path, + action="cleanup", + ) + + @patch('gaiaflow.cli.commands.minikube.load_imports') + def test_all_commands_handle_missing_project_gracefully(self, mock_load_imports): + mock_load_imports.return_value = self.mock_imports + self.mock_imports.gaiaflow_path_exists_in_state.return_value = False + + commands_and_args = [ + ["start", "--path", str(self.test_project_path)], + ["stop", "--path", str(self.test_project_path)], + ["restart", "--path", str(self.test_project_path)], + ["dockerize", "--path", str(self.test_project_path)], + ["create-config", "--path", str(self.test_project_path)], + ["create-secret", "--path", str(self.test_project_path), + "--name", "test", "--data", "key=value"], + ["cleanup", "--path", str(self.test_project_path)], + ] + + for command_args in commands_and_args: + with self.subTest(command=command_args[0]): + self.mock_imports.MinikubeManager.run.reset_mock() + + result = self.runner.invoke(prod_app, command_args) + + self.assertEqual(result.exit_code, 0) + self.assertIn("Please create a project with Gaiaflow", result.output) + + self.mock_imports.MinikubeManager.run.assert_not_called() + + def test_load_imports_function(self): + imports = load_imports() + + expected_attributes = [ + 'BaseAction', 'ExtendedAction', 'MinikubeManager', + 'create_gaiaflow_context_path', 'gaiaflow_path_exists_in_state', + 'parse_key_value_pairs' + ] + + for attr in expected_attributes: + with self.subTest(attribute=attr): + self.assertTrue(hasattr(imports, attr), + f"Missing attribute: {attr}") + + @patch('gaiaflow.cli.commands.minikube.load_imports') + def test_action_objects_comparison(self, mock_load_imports): + mock_load_imports.return_value = self.mock_imports + + result = self.runner.invoke(prod_app, [ + "start", + "--path", str(self.test_project_path) + ]) + + self.assertEqual(result.exit_code, 0) + + call_args = self.mock_imports.MinikubeManager.run.call_args + passed_action = call_args.kwargs['action'] + + self.assertEqual(passed_action, BaseAction.START) + self.assertEqual(passed_action.name, "start") diff --git a/tests/cli/commands/test_mlops.py b/tests/cli/commands/test_mlops.py new file mode 100644 index 0000000..b95126f --- /dev/null +++ b/tests/cli/commands/test_mlops.py @@ -0,0 +1,361 @@ +import unittest +from unittest.mock import patch, MagicMock, call +from pathlib import Path +from types import SimpleNamespace +from typer.testing import CliRunner + +from gaiaflow.constants import Service, DEFAULT_IMAGE_NAME, BaseAction, \ + ExtendedAction +from gaiaflow.cli.commands.mlops import app, load_imports + +class TestGaiaflowCLI(unittest.TestCase): + def setUp(self): + self.runner = CliRunner() + self.test_project_path = Path("/test/project") + self.test_gaiaflow_path = Path("/tmp/gaiaflow/test") + + self.mock_imports = SimpleNamespace( + BaseAction=BaseAction, + ExtendedAction=ExtendedAction, + MlopsManager=MagicMock(), + MinikubeManager=MagicMock(), + create_gaiaflow_context_path=MagicMock( + return_value=(self.test_gaiaflow_path, self.test_project_path) + ), + gaiaflow_path_exists_in_state=MagicMock(return_value=True), + save_project_state=MagicMock(), + ) + + @patch('gaiaflow.cli.commands.mlops.load_imports') + def test_start_command_with_all_services(self, mock_load_imports): + mock_load_imports.return_value = self.mock_imports + + result = self.runner.invoke(app, [ + "start", + "--path", str(self.test_project_path) + ]) + + self.assertEqual(result.exit_code, 0) + + self.mock_imports.create_gaiaflow_context_path.assert_called_once_with( + self.test_project_path + ) + self.mock_imports.gaiaflow_path_exists_in_state.assert_called_once_with( + self.test_gaiaflow_path, True + ) + + self.mock_imports.MlopsManager.run.assert_called_once_with( + gaiaflow_path=self.test_gaiaflow_path, + user_project_path=self.test_project_path, + force_new=False, + action=BaseAction.START, + service=Service.all, + cache=False, + jupyter_port=8895, + docker_build=False, + user_env_name=None, + env_tool="mamba", + ) + + @patch('gaiaflow.cli.commands.mlops.load_imports') + def test_start_command_with_specific_services(self, mock_load_imports): + mock_load_imports.return_value = self.mock_imports + + result = self.runner.invoke(app, [ + "start", + "--path", str(self.test_project_path), + "--service", "jupyter", + "--service", "airflow", + "--cache", + "--jupyter-port", "9000", + "--docker-build", + "--env", "myenv", + "--env-tool", "conda" + ]) + + self.assertEqual(result.exit_code, 0) + + self.assertEqual(self.mock_imports.MlopsManager.run.call_count, 2) + + expected_calls = [ + call( + gaiaflow_path=self.test_gaiaflow_path, + user_project_path=self.test_project_path, + force_new=False, + action=BaseAction.START, + service="jupyter", + cache=True, + jupyter_port=9000, + docker_build=True, + user_env_name="myenv", + env_tool="conda", + ), + call( + gaiaflow_path=self.test_gaiaflow_path, + user_project_path=self.test_project_path, + force_new=False, + action=BaseAction.START, + service="airflow", + cache=True, + jupyter_port=9000, + docker_build=True, + user_env_name="myenv", + env_tool="conda", + ) + ] + self.mock_imports.MlopsManager.run.assert_has_calls(expected_calls) + + @patch('gaiaflow.cli.commands.mlops.load_imports') + def test_start_command_saves_project_state_when_not_exists(self, mock_load_imports): + mock_load_imports.return_value = self.mock_imports + self.mock_imports.gaiaflow_path_exists_in_state.return_value = False + + result = self.runner.invoke(app, [ + "start", + "--path", str(self.test_project_path) + ]) + + self.assertEqual(result.exit_code, 0) + self.mock_imports.save_project_state.assert_called_once_with( + self.test_project_path, self.test_gaiaflow_path + ) + + @patch('gaiaflow.cli.commands.mlops.load_imports') + def test_stop_command_with_all_services(self, mock_load_imports): + mock_load_imports.return_value = self.mock_imports + + result = self.runner.invoke(app, [ + "stop", + "--path", str(self.test_project_path), + "--delete-volume" + ]) + + self.assertEqual(result.exit_code, 0) + + self.mock_imports.MlopsManager.run.assert_called_once_with( + gaiaflow_path=Path(self.test_gaiaflow_path), + user_project_path=Path(self.test_project_path), + service=Service.all, + action=BaseAction.STOP, + delete_volume=True, + ) + + @patch('gaiaflow.cli.commands.mlops.load_imports') + def test_stop_command_with_specific_services(self, mock_load_imports): + mock_load_imports.return_value = self.mock_imports + + result = self.runner.invoke(app, [ + "stop", + "--path", str(self.test_project_path), + "--service", "jupyter" + ]) + + self.assertEqual(result.exit_code, 0) + + self.mock_imports.MlopsManager.run.assert_called_once_with( + gaiaflow_path=Path(self.test_gaiaflow_path), + user_project_path=Path(self.test_project_path), + action=BaseAction.STOP, + service="jupyter", + delete_volume=False, + ) + + @patch('gaiaflow.cli.commands.mlops.load_imports') + def test_restart_command_with_all_services(self, mock_load_imports): + mock_load_imports.return_value = self.mock_imports + + result = self.runner.invoke(app, [ + "restart", + "--path", str(self.test_project_path), + "--force-new", + "--cache", + "--jupyter-port", "9001", + "--docker-build", + "--delete-volume" + ]) + + self.assertEqual(result.exit_code, 0) + + self.mock_imports.MlopsManager.run.assert_called_once_with( + gaiaflow_path=Path(self.test_gaiaflow_path), + user_project_path=Path(self.test_project_path), + force_new=True, + action=BaseAction.RESTART, + cache=True, + jupyter_port=9001, + delete_volume=True, + docker_build=True, + service=Service.all, + ) + + @patch("gaiaflow.cli.commands.mlops.load_imports") + def test_restart_command_with_specific_services(self, mock_load_imports): + mock_load_imports.return_value = self.mock_imports + + result = self.runner.invoke( + app, + [ + "restart", + "--path", + str(self.test_project_path), + "--service", + "jupyter", + "--service", + "airflow", + ], + ) + + self.assertEqual(result.exit_code, 0) + + self.assertEqual(self.mock_imports.MlopsManager.run.call_count, 2) + + expected_calls = [ + call( + gaiaflow_path=self.test_gaiaflow_path, + user_project_path=self.test_project_path, + force_new=False, + action=BaseAction.RESTART, + service="jupyter", + cache=False, + jupyter_port=8895, + delete_volume=False, + docker_build=False, + ), + call( + gaiaflow_path=self.test_gaiaflow_path, + user_project_path=self.test_project_path, + force_new=False, + action=BaseAction.RESTART, + service="airflow", + cache=False, + jupyter_port=8895, + delete_volume=False, + docker_build=False, + ), + ] + self.mock_imports.MlopsManager.run.assert_has_calls(expected_calls) + + @patch('gaiaflow.cli.commands.mlops.load_imports') + def test_cleanup_command(self, mock_load_imports): + mock_load_imports.return_value = self.mock_imports + + result = self.runner.invoke(app, [ + "cleanup", + "--path", str(self.test_project_path), + "--prune" + ]) + + self.assertEqual(result.exit_code, 0) + + self.mock_imports.MlopsManager.run.assert_called_once_with( + gaiaflow_path=Path(self.test_gaiaflow_path), + user_project_path=Path(self.test_project_path), + action=BaseAction.CLEANUP, + prune=True, + ) + + @patch('gaiaflow.cli.commands.mlops.load_imports') + def test_dockerize_command(self, mock_load_imports): + mock_load_imports.return_value = self.mock_imports + self.mock_imports.gaiaflow_path_exists_in_state.return_value = False + + result = self.runner.invoke(app, [ + "dockerize", + "--path", str(self.test_project_path), + "--image-name", "custom-image" + ]) + + self.assertEqual(result.exit_code, 0) + + self.mock_imports.save_project_state.assert_called_once_with( + self.test_project_path, self.test_gaiaflow_path + ) + + self.mock_imports.MinikubeManager.run.assert_called_once_with( + gaiaflow_path=self.test_gaiaflow_path, + user_project_path=self.test_project_path, + action=ExtendedAction.DOCKERIZE, + local=True, + image_name="custom-image" + ) + + @patch('gaiaflow.cli.commands.mlops.load_imports') + def test_dockerize_command_with_default_image_name(self, mock_load_imports): + mock_load_imports.return_value = self.mock_imports + + result = self.runner.invoke(app, [ + "dockerize", + "--path", str(self.test_project_path) + ]) + + self.assertEqual(result.exit_code, 0) + + self.mock_imports.MinikubeManager.run.assert_called_once() + call_args = self.mock_imports.MinikubeManager.run.call_args + self.assertEqual(call_args.kwargs['image_name'], DEFAULT_IMAGE_NAME) + + @patch('gaiaflow.cli.commands.mlops.load_imports') + def test_update_deps_command(self, mock_load_imports): + mock_load_imports.return_value = self.mock_imports + + result = self.runner.invoke( + app, ["update-deps", "--path", str(self.test_project_path)] + ) + + self.assertEqual(result.exit_code, 0) + + self.mock_imports.MlopsManager.run.assert_called_once_with( + gaiaflow_path=self.test_gaiaflow_path, + user_project_path=self.test_project_path, + action=ExtendedAction.UPDATE_DEPS, + ) + + @patch('gaiaflow.cli.commands.mlops.load_imports') + def test_commands_handle_missing_project_gracefully(self, mock_load_imports): + mock_load_imports.return_value = self.mock_imports + self.mock_imports.gaiaflow_path_exists_in_state.return_value = False + + failing_commands = ["stop", "restart", "cleanup"] + + for command in failing_commands: + with self.subTest(command=command): + result = self.runner.invoke(app, [ + command, + "--path", str(self.test_project_path) + ]) + + self.assertEqual(result.exit_code, 0) + self.assertIn("Please create a project with Gaiaflow", result.output) + + def test_load_imports_function(self): + imports = load_imports() + + expected_attributes = [ + 'BaseAction', 'ExtendedAction', 'MlopsManager', 'MinikubeManager', + 'create_gaiaflow_context_path', 'gaiaflow_path_exists_in_state', + 'save_project_state' + ] + + for attr in expected_attributes: + with self.subTest(attribute=attr): + self.assertTrue(hasattr(imports, attr), + f"Missing attribute: {attr}") + + @patch('gaiaflow.cli.commands.mlops.load_imports') + def test_argument_type_conversion(self, mock_load_imports): + mock_load_imports.return_value = self.mock_imports + + result = self.runner.invoke(app, [ + "start", + "--path", "/some/string/path" + ]) + + self.assertEqual(result.exit_code, 0) + + call_args = self.mock_imports.create_gaiaflow_context_path.call_args + self.assertIsInstance(call_args[0][0], Path) + self.assertEqual(str(call_args[0][0]), "/some/string/path") + + +if __name__ == '__main__': + unittest.main() \ No newline at end of file diff --git a/tests/cli/test_cli.py b/tests/cli/test_cli.py new file mode 100644 index 0000000..ee12b8a --- /dev/null +++ b/tests/cli/test_cli.py @@ -0,0 +1,16 @@ +import unittest +from typer.testing import CliRunner + +from gaiaflow.cli.cli import app as root_app + +runner = CliRunner() + + +class TestGaiaflowRootCLI(unittest.TestCase): + + def test_help_message_shows(self): + result = runner.invoke(root_app, ["--help"]) + self.assertEqual(result.exit_code, 0) + self.assertIn("Gaiaflow CLI is a manager tool", result.stdout) + self.assertIn("dev", result.stdout) + self.assertIn("prod-local", result.stdout) \ No newline at end of file diff --git a/tests/core/test_create_task.py b/tests/core/test_create_task.py new file mode 100644 index 0000000..674a054 --- /dev/null +++ b/tests/core/test_create_task.py @@ -0,0 +1,147 @@ +import unittest +from unittest.mock import Mock, patch + +from gaiaflow.core.create_task import GaiaflowMode, create_task + + +class TestCreateTask(unittest.TestCase): + def setUp(self): + self.mock_operator_instance = Mock() + self.mock_operator_instance.create_task.return_value = "mock_task" + + self.mock_operator_class = Mock() + self.mock_operator_class.return_value = self.mock_operator_instance + + self.operator_map_patcher = patch("gaiaflow.core.create_task.OPERATOR_MAP") + self.mock_operator_map = self.operator_map_patcher.start() + self.mock_operator_map.get.return_value = self.mock_operator_class + + def tearDown(self): + self.operator_map_patcher.stop() + + def test_enum_values(self): + self.assertEqual(GaiaflowMode.DEV.value, "dev") + self.assertEqual(GaiaflowMode.DEV_DOCKER.value, "dev_docker") + self.assertEqual(GaiaflowMode.PROD_LOCAL.value, "prod_local") + self.assertEqual(GaiaflowMode.PROD.value, "prod") + + def test_enum_invalid_value(self): + with self.assertRaises(ValueError): + GaiaflowMode("invalid_mode") + + def test_enum_creation_from_string(self): + self.assertEqual(GaiaflowMode("dev"), GaiaflowMode.DEV) + self.assertEqual(GaiaflowMode("dev_docker"), GaiaflowMode.DEV_DOCKER) + self.assertEqual(GaiaflowMode("prod_local"), GaiaflowMode.PROD_LOCAL) + self.assertEqual(GaiaflowMode("prod"), GaiaflowMode.PROD) + + def test_create_task_basic(self): + result = create_task(task_id="test_task", func_path="test.module:function") + + self.assertEqual(result, "mock_task") + self.mock_operator_class.assert_called_once() + self.mock_operator_instance.create_task.assert_called_once() + + def test_create_task_with_all_parameters(self): + mock_dag = Mock() + mock_dag.params = {"param1": "value1"} + + result = create_task( + task_id="test_task", + func_path="test.module:function", + func_kwargs={"key": "value"}, + func_args=["arg1", "arg2"], + image="test_image:latest", + mode="prod", + secrets=["secret1", "secret2"], + env_vars={"ENV_VAR": "value"}, + retries=5, + dag=mock_dag, + ) + + self.assertEqual(result, "mock_task") + + self.mock_operator_class.assert_called_once_with( + task_id="test_task", + func_path="test.module:function", + func_args=["arg1", "arg2"], + func_kwargs={"key": "value"}, + image="test_image:latest", + secrets=["secret1", "secret2"], + env_vars={"ENV_VAR": "value"}, + retries=5, + params={"param1": "value1"}, + mode=GaiaflowMode.PROD, + ) + + def test_create_task_default_values(self): + create_task(task_id="test_task", func_path="test.module:function") + + call_args = self.mock_operator_class.call_args + self.assertEqual(call_args.kwargs["func_args"], []) + self.assertEqual(call_args.kwargs["func_kwargs"], {}) + self.assertEqual(call_args.kwargs["env_vars"], {}) + self.assertEqual(call_args.kwargs["retries"], 3) + self.assertEqual(call_args.kwargs["mode"], GaiaflowMode.DEV) + self.assertEqual(call_args.kwargs["params"], {}) + + def test_create_task_invalid_mode(self): + with self.assertRaises(ValueError) as context: + create_task( + task_id="test_task", + func_path="test.module:function", + mode="invalid_mode", + ) + + self.assertIn("env must be one of", str(context.exception)) + self.assertIn("invalid_mode", str(context.exception)) + + def test_create_task_no_operator_for_mode(self): + self.mock_operator_map.get.return_value = None + + with self.assertRaises(ValueError) as context: + create_task( + task_id="test_task", func_path="test.module:function", mode="dev" + ) + + self.assertIn("No task creation operator defined for", str(context.exception)) + + def test_create_task_with_dag_no_params(self): + mock_dag = Mock(spec=[]) # DAG without params attribute + + create_task(task_id="test_task", func_path="test.module:function", dag=mock_dag) + + call_args = self.mock_operator_class.call_args + self.assertEqual(call_args.kwargs["params"], {}) + + def test_create_task_none_values_converted_to_defaults(self): + create_task( + task_id="test_task", + func_path="test.module:function", + func_kwargs=None, + func_args=None, + env_vars=None, + ) + + call_args = self.mock_operator_class.call_args + self.assertEqual(call_args.kwargs["func_args"], []) + self.assertEqual(call_args.kwargs["func_kwargs"], {}) + self.assertEqual(call_args.kwargs["env_vars"], {}) + + def test_create_task_different_modes(self): + modes = ["dev", "dev_docker", "prod_local", "prod"] + expected_enums = [ + GaiaflowMode.DEV, + GaiaflowMode.DEV_DOCKER, + GaiaflowMode.PROD_LOCAL, + GaiaflowMode.PROD, + ] + + for mode, expected_enum in zip(modes, expected_enums): + with self.subTest(mode=mode): + create_task( + task_id="test_task", func_path="test.module.function", mode=mode + ) + + call_args = self.mock_operator_class.call_args + self.assertEqual(call_args.kwargs["mode"], expected_enum) diff --git a/tests/core/test_operators.py b/tests/core/test_operators.py new file mode 100644 index 0000000..12f9414 --- /dev/null +++ b/tests/core/test_operators.py @@ -0,0 +1,395 @@ +import json +import unittest +from unittest.mock import patch + +import gaiaflow.core.operators as operators +from gaiaflow.core.create_task import GaiaflowMode + + +class TestBaseTaskOperator(unittest.TestCase): + def setUp(self): + self.base_op = operators.BaseTaskOperator( + task_id="t1", + func_path="mymod:func", + func_args=[1, operators.FromTask("taskX", "key1")], + func_kwargs={"foo": operators.FromTask("taskY")}, + image="img", + secrets=["s1"], + env_vars={"ENV": "value"}, + retries=1, + params={"p": "q"}, + mode="dev", + ) + + def test_resolve_xcom_value_default_key(self): + val = self.base_op._resolve_xcom_value({"task": "up_task"}) + self.assertIn("ti.xcom_pull(task_ids='up_task')", val) + + def test_resolve_xcom_value_custom_key(self): + val = self.base_op._resolve_xcom_value({"task": "up_task", "key": "out"}) + self.assertIn("['out']", val) + + def test_resolve_args_kwargs(self): + args, kwargs = self.base_op.resolve_args_kwargs() + self.assertEqual(args[1], ("{{ ti.xcom_pull(task_ids='taskX')['key1'] }}")) + self.assertEqual(kwargs["foo"], "{{ ti.xcom_pull(task_ids='taskY') }}") + + def test_create_func_env_vars(self): + envs = self.base_op.create_func_env_vars() + self.assertIn("FUNC_PATH", envs) + self.assertTrue(isinstance(json.loads(envs["FUNC_ARGS"]), list)) + self.assertTrue(isinstance(json.loads(envs["FUNC_KWARGS"]), dict)) + + def test_create_task(self): + with self.assertRaises(NotImplementedError): + self.base_op.create_task() + + +class TestOperators(unittest.TestCase): + def test_to_dict(self): + ft = operators.FromTask("upstream_task", key="mykey") + self.assertEqual(ft.to_dict(), {"task": "upstream_task", "key": "mykey"}) + + def test_split_args_and_kwargs(self): + ft1 = operators.FromTask("task1", "key1") + ft2 = operators.FromTask("task1") + args, x_args, kwargs, x_kwargs = operators.split_args_kwargs( + func_args=[1, ft1, "a"], func_kwargs={"x": 1, "key2": ft2} + ) + self.assertEqual(args, [1, "a"]) + self.assertEqual(kwargs, {"x": 1}) + self.assertEqual(x_args, {"1": {"task": "task1", "key": "key1"}}) + self.assertEqual(x_kwargs, {"key2": {"task": "task1", "key": "return_value"}}) + + @patch("gaiaflow.core.operators.ExternalPythonOperator") + def test_create_dev_task(self, mock_ext_op): + op = operators.DevTaskOperator( + task_id="devtask", + func_path="mod:fn", + func_args=[1], + func_kwargs={"test": "123"}, + image=None, + secrets=[], + env_vars={}, + retries=2, + params={"x": "y"}, + mode=GaiaflowMode.DEV, + ) + op.create_task() + args, kwargs = mock_ext_op.call_args + self.assertEqual(kwargs["op_kwargs"]["func_path"], "mod:fn") + self.assertEqual(kwargs["op_kwargs"]["args"], [1]) + self.assertEqual( + kwargs["op_kwargs"]["kwargs"], {"params": {"x": "y"}, "test": "123"} + ) + self.assertEqual(kwargs["expect_airflow"], False) + self.assertEqual(kwargs["expect_pendulum"], False) + + def test_create_prod_local_task_no_image(self): + with self.assertRaises(ValueError): + operators.ProdLocalTaskOperator( + task_id="prodlocaltask", + func_path="mod:fn", + func_args=[1], + func_kwargs={"test": "123"}, + image=None, + secrets=[], + env_vars={}, + retries=2, + params={"x": "y"}, + mode=GaiaflowMode.PROD_LOCAL, + ).create_task() + + @patch("gaiaflow.core.operators.KubernetesPodOperator") + def test_create_prod_local_task(self, mock_ext_op): + op = operators.ProdLocalTaskOperator( + task_id="prodlocaltask", + func_path="mod:fn", + func_args=[1], + func_kwargs={"test": "123"}, + image="random_image:v1", + secrets=[], + env_vars={}, + retries=2, + params={"x": "y"}, + mode=GaiaflowMode.PROD_LOCAL, + ) + op.create_task() + args, kwargs = mock_ext_op.call_args + self.assertEqual(kwargs["image"], "random_image:v1") + self.assertEqual(kwargs["cmds"], ["python", "-m", "runner"]) + self.assertEqual( + kwargs["env_vars"], + { + "MODE": "prod_local", + "PARAMS_X": "{{ params.x }}", + "MLFLOW_TRACKING_URI": "http://192.168.49.1:5000", + "MLFLOW_S3_ENDPOINT_URL": "http://192.168.49.1:9000", + "AWS_ACCESS_KEY_ID": "minio", + "AWS_SECRET_ACCESS_KEY": "minio123", + "FUNC_PATH": "mod:fn", + "FUNC_ARGS": "[1]", + "FUNC_KWARGS": '{"test": "123"}', + }, + ) + self.assertEqual(kwargs["params"], {"x": "y"}) + self.assertEqual(kwargs["in_cluster"], False) + + with patch("platform.system", return_value="Windows"): + op.create_task() + args, kwargs = mock_ext_op.call_args + + self.assertEqual(kwargs["image"], "random_image:v1") + self.assertEqual(kwargs["cmds"], ["python", "-m", "runner"]) + self.assertEqual( + kwargs["env_vars"], + { + "MODE": "prod_local", + "PARAMS_X": "{{ params.x }}", + "MLFLOW_TRACKING_URI": "http://host.docker.internal:5000", + "MLFLOW_S3_ENDPOINT_URL": "http://host.docker.internal:9000", + "AWS_ACCESS_KEY_ID": "minio", + "AWS_SECRET_ACCESS_KEY": "minio123", + "FUNC_PATH": "mod:fn", + "FUNC_ARGS": "[1]", + "FUNC_KWARGS": '{"test": "123"}', + }, + ) + self.assertEqual(kwargs["params"], {"x": "y"}) + self.assertEqual(kwargs["in_cluster"], False) + + def test_create_task_unknown_profile_raises(self): + op = operators.ProdLocalTaskOperator( + task_id="prodlocaltask", + func_path="mod:fn", + func_args=[], + func_kwargs={}, + image="img", + secrets=[], + env_vars={}, + retries=1, + params={"resource_profile": "unknown"}, + mode=GaiaflowMode.PROD_LOCAL, + ) + with self.assertRaises(ValueError): + op.create_task() + + @patch("gaiaflow.core.operators.KubernetesPodOperator") + def test_create_prod_task(self, mock_ext_op): + op = operators.ProdTaskOperator( + task_id="prodtask", + func_path="mod:fn", + func_args=[1], + func_kwargs={"test": "123"}, + image="random_image:v1", + secrets=["my-secret"], + env_vars={}, + retries=2, + params={"x": "y"}, + mode=GaiaflowMode.PROD, + ) + op.create_task() + args, kwargs = mock_ext_op.call_args + print(args, kwargs) + self.assertEqual(kwargs["image"], "random_image:v1") + self.assertEqual(kwargs["cmds"], ["python", "-m", "runner"]) + self.assertEqual( + kwargs["env_vars"], + { + "MODE": "prod", + "PARAMS_X": "{{ params.x }}", + "MLFLOW_TRACKING_URI": "http://192.168.49.1:5000", + "MLFLOW_S3_ENDPOINT_URL": "http://192.168.49.1:9000", + "AWS_ACCESS_KEY_ID": "minio", + "AWS_SECRET_ACCESS_KEY": "minio123", + "FUNC_PATH": "mod:fn", + "FUNC_ARGS": "[1]", + "FUNC_KWARGS": '{"test": "123"}', + }, + ) + self.assertEqual(kwargs["params"], {"x": "y"}) + self.assertEqual(kwargs["in_cluster"], True) + self.assertEqual( + kwargs["env_from"][0].to_dict(), + { + "config_map_ref": None, + "prefix": None, + "secret_ref": {"name": "my-secret", "namespace": None}, + }, + ) + # self.assertEqual( + # kwargs["container_resources"].to_dict(), + # { + # "claims": None, + # "limits": {"cpu": "500m", "memory": "1Gi"}, + # "requests": {"cpu": "250m", "memory": "512Mi"}, + # }, + # ) + + @patch("gaiaflow.core.operators.KubernetesPodOperator") + def test_create_prod_task_windows(self, mock_ext_op): + op = operators.ProdTaskOperator( + task_id="prodtask", + func_path="mod:fn", + func_args=[1], + func_kwargs={"test": "123"}, + image="random_image:v1", + secrets=["my-secret"], + env_vars={}, + retries=2, + params={"x": "y"}, + mode=GaiaflowMode.PROD, + ) + with patch("platform.system", return_value="Windows"): + op.create_task() + args, kwargs = mock_ext_op.call_args + + self.assertEqual(kwargs["image"], "random_image:v1") + self.assertEqual(kwargs["cmds"], ["python", "-m", "runner"]) + self.assertEqual( + kwargs["env_vars"], + { + "MODE": "prod", + "PARAMS_X": "{{ params.x }}", + "MLFLOW_TRACKING_URI": "http://host.docker.internal:5000", + "MLFLOW_S3_ENDPOINT_URL": "http://host.docker.internal:9000", + "AWS_ACCESS_KEY_ID": "minio", + "AWS_SECRET_ACCESS_KEY": "minio123", + "FUNC_PATH": "mod:fn", + "FUNC_ARGS": "[1]", + "FUNC_KWARGS": '{"test": "123"}', + }, + ) + self.assertEqual(kwargs["params"], {"x": "y"}) + self.assertEqual(kwargs["in_cluster"], True) + + @patch("gaiaflow.core.operators.KubernetesPodOperator") + def test_create_prod_task_custom_env_vars(self, mock_ext_op): + op = operators.ProdTaskOperator( + task_id="prodtask", + func_path="mod:fn", + func_args=[1], + func_kwargs={"test": "123"}, + image="random_image:v1", + secrets=["my-secret"], + env_vars={"AWS_ACCESS_KEY_ID": "test", "AWS_SECRET_ACCESS_KEY": + "test2", "MINIKUBE_GATEWAY": "test3"}, + retries=2, + params={"x": "y"}, + mode=GaiaflowMode.PROD, + ) + op.create_task() + args, kwargs = mock_ext_op.call_args + print(args, kwargs) + self.assertEqual(kwargs["image"], "random_image:v1") + self.assertEqual(kwargs["cmds"], ["python", "-m", "runner"]) + self.assertEqual( + kwargs["env_vars"], + { + "MODE": "prod", + "PARAMS_X": "{{ params.x }}", + "MLFLOW_TRACKING_URI": "http://test3:5000", + "MLFLOW_S3_ENDPOINT_URL": "http://test3:9000", + "AWS_ACCESS_KEY_ID": "test", + "AWS_SECRET_ACCESS_KEY": "test2", + "MINIKUBE_GATEWAY": "test3", + "FUNC_PATH": "mod:fn", + "FUNC_ARGS": "[1]", + "FUNC_KWARGS": '{"test": "123"}', + }, + ) + self.assertEqual(kwargs["params"], {"x": "y"}) + self.assertEqual(kwargs["in_cluster"], True) + self.assertEqual( + kwargs["env_from"][0].to_dict(), + { + "config_map_ref": None, + "prefix": None, + "secret_ref": {"name": "my-secret", "namespace": None}, + }, + ) + # self.assertEqual( + # kwargs["container_resources"].to_dict(), + # { + # "claims": None, + # "limits": {"cpu": "500m", "memory": "1Gi"}, + # "requests": {"cpu": "250m", "memory": "512Mi"}, + # }, + # ) + + @patch("gaiaflow.core.operators.DockerOperator") + def test_create_dev_docker_task(self, mock_ext_op): + op = operators.DockerTaskOperator( + task_id="devdockertask", + func_path="mod:fn", + func_args=[1], + func_kwargs={"test": "123"}, + image="random_image:v1", + secrets=[], + env_vars={}, + retries=2, + params={"x": "y"}, + mode=GaiaflowMode.DEV_DOCKER, + ) + op.create_task() + args, kwargs = mock_ext_op.call_args + self.assertEqual(kwargs["image"], "random_image:v1") + self.assertEqual(kwargs["command"], ["python", "-m", "runner"]) + self.assertEqual(kwargs["docker_url"], "unix://var/run/docker.sock") + self.assertEqual(kwargs["retrieve_output"], True) + self.assertEqual(kwargs["retrieve_output_path"], "/tmp/script.out") + self.assertEqual( + kwargs["environment"], + { + "MODE": "dev_docker", + "PARAMS_X": "{{ params.x }}", + "MLFLOW_TRACKING_URI": "http://mlflow:5000", + "MLFLOW_S3_ENDPOINT_URL": "http://minio:9000", + "AWS_ACCESS_KEY_ID": "minio", + "AWS_SECRET_ACCESS_KEY": "minio123", + "FUNC_PATH": "mod:fn", + "FUNC_ARGS": "[1]", + "FUNC_KWARGS": '{"test": "123"}', + }, + ) + + @patch("gaiaflow.core.operators.DockerOperator") + def test_create_dev_docker_task_with_custom_env_vars(self, mock_ext_op): + op = operators.DockerTaskOperator( + task_id="devdockertask", + func_path="mod:fn", + func_args=[1], + func_kwargs={"test": "123"}, + image="random_image:v1", + secrets=[], + env_vars={ + "MLFLOW_TRACKING_URI": "test", + "MLFLOW_S3_ENDPOINT_URL": "test2", + "AWS_ACCESS_KEY_ID": "test3", + "AWS_SECRET_ACCESS_KEY": "test4" + }, + retries=2, + params={"x": "y"}, + mode=GaiaflowMode.DEV_DOCKER, + ) + op.create_task() + args, kwargs = mock_ext_op.call_args + self.assertEqual(kwargs["image"], "random_image:v1") + self.assertEqual(kwargs["command"], ["python", "-m", "runner"]) + self.assertEqual(kwargs["docker_url"], "unix://var/run/docker.sock") + self.assertEqual(kwargs["retrieve_output"], True) + self.assertEqual(kwargs["retrieve_output_path"], "/tmp/script.out") + self.assertEqual( + kwargs["environment"], + { + "MODE": "dev_docker", + "PARAMS_X": "{{ params.x }}", + "MLFLOW_TRACKING_URI": "test", + "MLFLOW_S3_ENDPOINT_URL": "test2", + "AWS_ACCESS_KEY_ID": "test3", + "AWS_SECRET_ACCESS_KEY": "test4", + "FUNC_PATH": "mod:fn", + "FUNC_ARGS": "[1]", + "FUNC_KWARGS": '{"test": "123"}', + }, + ) diff --git a/tests/core/test_runner.py b/tests/core/test_runner.py new file mode 100644 index 0000000..90f1eae --- /dev/null +++ b/tests/core/test_runner.py @@ -0,0 +1,183 @@ +import os +import shutil +import tempfile +import types +import unittest +from unittest.mock import Mock, mock_open, patch + +from gaiaflow.core import runner +from gaiaflow.testing import set_env_cm + + +class TestRunner(unittest.TestCase): + def setUp(self): + self.temp_dir = tempfile.mkdtemp() + self.xcom_dir = os.path.join(self.temp_dir, "airflow", "xcom") + + def tearDown(self): + shutil.rmtree(self.temp_dir, ignore_errors=True) + + def test_extract_params_empty_env(self): + with set_env_cm(): + result = runner._extract_params_from_env() + self.assertEqual(result, {}) + + def test_extract_params_with_params(self): + with set_env_cm( + PARAMS_KEY1="value1", + PARAMS_KEY2="value2", + OTHER_VAR="not_included", + PARAMS_NESTED_KEY="nested_value", + ): + result = runner._extract_params_from_env() + expected = {"key1": "value1", "key2": "value2", "nested_key": "nested_value"} + self.assertEqual(result, expected) + + def test_extract_params_case_conversion(self): + with set_env_cm( + PARAMS_UPPER_CASE="value", + PARAMS_MixedCase="value2", + ): + result = runner._extract_params_from_env() + expected = {"upper_case": "value", "mixedcase": "value2"} + self.assertEqual(result, expected) + + def test_extract_params_custom_prefix(self): + with set_env_cm( + CUSTOM_KEY1="value", + CUSTOM_KEY2="value2", + ): + result = runner._extract_params_from_env("CUSTOM_") + expected = {"key1": "value", "key2": "value2"} + self.assertEqual(result, expected) + + @patch("gaiaflow.core.runner.os.makedirs") + @patch("gaiaflow.core.runner.open", new_callable=mock_open) + @patch("gaiaflow.core.runner.json.dump") + def test_write_xcom_result_success( + self, mock_json_dump, mock_file_open, mock_makedirs + ): + result = {"key": "value", "number": 123} + + runner._write_xcom_result(result) + + mock_makedirs.assert_called_once_with("/airflow/xcom", exist_ok=True) + mock_file_open.assert_called_once_with("/airflow/xcom/return.json", "w") + mock_json_dump.assert_called_once_with( + result, mock_file_open.return_value.__enter__.return_value + ) + + @patch("gaiaflow.core.runner.os.makedirs") + def test_write_xcom_result_makedirs_failure(self, mock_makedirs): + result = {"key": "value"} + mock_makedirs.side_effect = OSError("Permission denied") + + with self.assertRaises(OSError): + runner._write_xcom_result(result) + + def test_run_no_func_path_raises_error(self): + with self.assertRaises(ValueError) as context: + runner.run() + + self.assertEqual(str(context.exception), "func_path must be provided") + + def test_import_function_success(self): + def dummy_func(): + return "ok" + + fake_module = types.SimpleNamespace(myfunc=dummy_func) + + with patch("importlib.import_module", return_value=fake_module): + func = runner._import_function("fake_module:myfunc") + + self.assertEqual(func, dummy_func) + + def test_import_function_invalid_path(self): + with self.assertRaises(ValueError): + runner._import_function("not_a_valid_path") + + def test_resolve_inputs_dev_mode(self): + func_path, args, kwargs = runner._resolve_inputs( + "mymod:func", [1, 2], {"a": 3}, "dev" + ) + self.assertEqual(func_path, "mymod:func") + self.assertEqual(args, [1, 2]) + self.assertEqual(kwargs, {"a": 3}) + + def test_resolve_inputs_prod_local_mode(self): + # This test will be the same for dev_docker and prod mode as well + # because they all expect the func_path, args and kwargs as env + # variables. + with set_env_cm( + FUNC_PATH="mod:func", + FUNC_ARGS="[10, 20]", + FUNC_KWARGS='{"foo": "bar"}', + ): + with patch( + "gaiaflow.core.runner._extract_params_from_env", return_value={"x": 1} + ): + func_path, args, kwargs = runner._resolve_inputs( + None, None, None, "prod" + ) + + self.assertEqual(func_path, "mod:func") + self.assertEqual(args, [10, 20]) + self.assertEqual(kwargs, {"foo": "bar", "params": {"x": 1}}) + + def test_run_dev_mode(self): + dummy_func = Mock(return_value={"res": "ok"}) + fake_module = types.SimpleNamespace(myfunc=dummy_func) + + with patch("importlib.import_module", return_value=fake_module): + result = runner.run( + func_path="fake_module:myfunc", + args=[1, 2], + kwargs={"k": "v"}, + ) + + self.assertEqual(result, {"res": "ok"}) + dummy_func.assert_called_once_with(1, 2, k="v") + + def test_run_prod_mode_with_xcom(self): + dummy_func = Mock(return_value={"done": True}) + fake_module = types.SimpleNamespace(myfunc=dummy_func) + + with set_env_cm( + MODE="prod", + FUNC_PATH="fake_module:myfunc", + FUNC_ARGS="[100]", + FUNC_KWARGS='{"alpha": 1}', + ): + with ( + patch("gaiaflow.core.runner._write_xcom_result") as mock_xcom, + patch("importlib.import_module", return_value=fake_module), + ): + result = runner.run() + + self.assertEqual(result, {"done": True}) + dummy_func.assert_called_once_with(100, alpha=1, params={}) + mock_xcom.assert_called_once_with({"done": True}) + + @patch("gaiaflow.core.runner.pickle.dump") + @patch("gaiaflow.core.runner.open", new_callable=mock_open) + def test_run_dev_docker_mode(self, mock_file_open, mock_pickle_dump): + dummy_func = Mock(return_value={"docker": "ok"}) + fake_module = types.SimpleNamespace(myfunc=dummy_func) + + with set_env_cm( + MODE="dev_docker", + FUNC_PATH="fake_module:myfunc", + FUNC_ARGS="[1]", + FUNC_KWARGS='{"flag": true}', + ): + with patch("importlib.import_module", return_value=fake_module): + result = runner.run() + + self.assertEqual(result, {"docker": "ok"}) + dummy_func.assert_called_once_with(1, flag=True, params={}) + + mock_file_open.assert_called_once_with("/tmp/script.out", "wb+") + mock_pickle_dump.assert_called_once_with( + {"docker": "ok"}, + mock_file_open.return_value.__enter__.return_value, + ) diff --git a/tests/core/test_utils.py b/tests/core/test_utils.py index 143b625..f9cec6a 100644 --- a/tests/core/test_utils.py +++ b/tests/core/test_utils.py @@ -1,12 +1,17 @@ +import subprocess import unittest +from unittest.mock import patch + from kubernetes.client import V1EnvFromSource, V1SecretReference -from gaiaflow.core.utils import inject_params_as_env_vars, \ - build_env_from_secrets +from gaiaflow.core.utils import ( + build_env_from_secrets, + docker_network_gateway, + inject_params_as_env_vars, +) class TestUtils(unittest.TestCase): - def test_inject_params_as_env_vars(self): params = {"foo": "bar", "baz": "qux"} expected = { @@ -39,3 +44,25 @@ def test_build_env_from_secrets(self): def test_build_env_from_secrets_empty(self): result = build_env_from_secrets([]) self.assertEqual(result, []) + + def test_gateway_found(self): + mock_output = '{"Gateway": "172.18.0.1"}' + with patch("gaiaflow.core.utils.subprocess.run") as mock_run: + mock_run.return_value = subprocess.CompletedProcess( + args=["docker"], returncode=0, stdout=mock_output + ) + result = docker_network_gateway() + self.assertEqual(result, "172.18.0.1") + + def test_gateway_not_found(self): + with patch("gaiaflow.core.utils.subprocess.run") as mock_run: + mock_run.return_value = subprocess.CompletedProcess( + args=["docker"], returncode=0, stdout="{}" + ) + result = docker_network_gateway() + self.assertIsNone(result) + + def test_docker_not_installed(self): + with patch("subprocess.run", side_effect=FileNotFoundError): + result = docker_network_gateway() + self.assertIsNone(result) diff --git a/tests/managers/__init__.py b/tests/managers/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/managers/test_base_manager.py b/tests/managers/test_base_manager.py new file mode 100644 index 0000000..6671aed --- /dev/null +++ b/tests/managers/test_base_manager.py @@ -0,0 +1,80 @@ +import unittest +from pathlib import Path +from unittest.mock import MagicMock + +from gaiaflow.constants import BaseAction +from gaiaflow.managers.base_manager import BaseGaiaflowManager + + +class DummyManager(BaseGaiaflowManager): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.started = False + self.stopped = False + self.cleaned = False + + def start(self): + self.started = True + + def stop(self): + self.stopped = True + + def cleanup(self): + self.cleaned = True + + +class TestBaseGaiaflowManager(unittest.TestCase): + def setUp(self): + self.gaiaflow_path = Path("/tmp/gaiaflow") + self.user_project_path = Path("/tmp/project") + + def test_invalid_action_raises_value_error(self): + class FakeAction: + name = "FAKE" + + with self.assertRaises(ValueError) as ctx: + DummyManager(self.gaiaflow_path, self.user_project_path, FakeAction()) + self.assertIn("Invalid action", str(ctx.exception)) + + def test_valid_action_initialization(self): + mgr = DummyManager( + self.gaiaflow_path, + self.user_project_path, + BaseAction.START, + force_new=True, + prune=True, + ) + self.assertEqual(mgr.gaiaflow_path, self.gaiaflow_path) + self.assertEqual(mgr.user_project_path, self.user_project_path) + self.assertEqual(mgr.action, BaseAction.START) + self.assertTrue(mgr.force_new) + self.assertTrue(mgr.prune) + + def test_restart_calls_stop_and_start(self): + mgr = DummyManager( + self.gaiaflow_path, self.user_project_path, BaseAction.RESTART + ) + mgr.stop = MagicMock() + mgr.start = MagicMock() + + mgr.restart() + + mgr.stop.assert_called_once() + mgr.start.assert_called_once() + + def test_get_valid_actions_default(self): + mgr = DummyManager( + self.gaiaflow_path, self.user_project_path, BaseAction.START + ) + actions = mgr._get_valid_actions() + self.assertIn(BaseAction.START, actions) + self.assertIn(BaseAction.STOP, actions) + self.assertIn(BaseAction.RESTART, actions) + self.assertIn(BaseAction.CLEANUP, actions) + + def test_cleanup_is_called(self): + mgr = DummyManager( + self.gaiaflow_path, self.user_project_path, BaseAction.CLEANUP + ) + mgr.cleanup() + self.assertTrue(mgr.cleaned) diff --git a/tests/managers/test_minikube_manager.py b/tests/managers/test_minikube_manager.py new file mode 100644 index 0000000..aeb6723 --- /dev/null +++ b/tests/managers/test_minikube_manager.py @@ -0,0 +1,435 @@ +import unittest +from pathlib import Path +import tempfile +from unittest.mock import patch, MagicMock +import subprocess + +import yaml + +from gaiaflow.constants import BaseAction, ExtendedAction +from gaiaflow.managers.minikube_manager import ( + MinikubeManager, + MinikubeHelper, + DockerHelper, + KubeConfigHelper, + temporary_copy, +) + + +class TestTemporaryCopy(unittest.TestCase): + def test_temporary_copy_creates_and_deletes(self): + with tempfile.TemporaryDirectory() as tmpdir: + src = Path(tmpdir) / "src.txt" + dest = Path(tmpdir) / "dest.txt" + + src.write_text("hello") + + with temporary_copy(src, dest): + self.assertTrue(dest.exists()) + self.assertEqual(dest.read_text(), "hello") + + self.assertFalse(dest.exists()) + +class TestMinikubeHelper(unittest.TestCase): + def setUp(self): + self.helper = MinikubeHelper(profile="test-profile") + + def test_profile_name_is_stored(self): + self.assertEqual(self.helper.profile, "test-profile") + + def test_has_expected_methods(self): + for method in ["is_running", "start", "stop", "cleanup", "run_cmd"]: + self.assertTrue(callable(getattr(self.helper, method))) + + @patch("subprocess.run") + def test_is_running_true(self, mock_run): + mock_run.return_value = MagicMock(stdout=b"Running") + self.assertTrue(self.helper.is_running()) + + @patch("subprocess.run") + def test_is_running_false(self, mock_run): + mock_run.return_value = MagicMock(stdout=b"Stopped") + self.assertFalse(self.helper.is_running()) + + @patch("gaiaflow.managers.minikube_manager.run") + @patch("gaiaflow.managers.minikube_manager.is_wsl", return_value=False) + @patch.object(MinikubeHelper, "is_running", return_value=False) + def test_start_success(self, mock_is_running, mock_is_wsl, mock_run): + self.helper.start() + mock_run.assert_called() + args = mock_run.call_args + self.assertEqual(args[0][0], ['minikube', 'start', '--profile', + 'test-profile', '--driver=docker', '--cpus=4', '--memory=4g']) + + @patch("gaiaflow.managers.minikube_manager.run") + @patch("gaiaflow.managers.minikube_manager.is_wsl", return_value=True) + @patch.object(MinikubeHelper, "is_running", return_value=False) + def test_start_success_wsl(self, mock_is_running, mock_is_wsl, mock_run): + self.helper.start() + mock_run.assert_called() + args = mock_run.call_args + self.assertEqual( + args[0][0], + [ + "minikube", + "start", + "--profile", + "test-profile", + "--driver=docker", + "--cpus=4", + "--memory=4g", + "--extra-config=kubelet.cgroup-driver=cgroupfs" + ], + ) + + @patch("gaiaflow.managers.minikube_manager.run") + @patch.object(MinikubeHelper, "is_running", return_value=True) + def test_start_already_running(self, mock_is_running, + mock_run): + self.helper.start() + mock_run.assert_not_called() + + @patch( + "gaiaflow.managers.minikube_manager.run", + side_effect=subprocess.CalledProcessError(1, "cmd"), + ) + @patch("gaiaflow.managers.minikube_manager.is_wsl", return_value=False) + @patch.object(MinikubeHelper, "is_running", return_value=False) + @patch.object(MinikubeHelper, "cleanup") + def test_start_retries_after_cleanup(self, mock_cleanup, *_): + with self.assertRaises(subprocess.CalledProcessError): + self.helper.start() + mock_cleanup.assert_called() + + @patch("gaiaflow.managers.minikube_manager.run") + def test_stop(self, mock_run): + self.helper.stop() + mock_run.assert_called() + + @patch("gaiaflow.managers.minikube_manager.run") + def test_cleanup(self, mock_run): + self.helper.cleanup() + mock_run.assert_called() + + @patch("subprocess.run") + def test_run_cmd(self, mock_run): + self.helper.run_cmd(["status"]) + mock_run.assert_called() + +class TestDockerHelper(unittest.TestCase): + def setUp(self): + self.tmpdir = tempfile.TemporaryDirectory() + self.project_path = Path(self.tmpdir.name) + self.helper = DockerHelper( + image_name="test-image", + project_path=self.project_path, + local=True, + minikube_helper=MinikubeHelper(), + ) + + def tearDown(self): + self.tmpdir.cleanup() + + def _write_dockerfile(self): + dockerfile = self.project_path / "Dockerfile" + dockerfile.write_text("ENV TEST=1\nENTRYPOINT test.sh\n") + return dockerfile + + def test_stores_init_params(self): + self.assertEqual(self.helper.image_name, "test-image") + self.assertEqual(self.helper.project_path, self.project_path) + self.assertTrue(self.helper.local) + + def test_has_expected_methods(self): + for method in [ + "build_image", + "_update_dockerfile", + "_build_local", + "_build_minikube", + ]: + self.assertTrue(hasattr(self.helper, method)) + + @patch( + "gaiaflow.managers.minikube_manager.find_python_packages", + return_value=["mypkg"], + ) + def test_update_dockerfile_inserts_copy(self, _): + dockerfile = self._write_dockerfile() + self.helper._update_dockerfile(dockerfile) + text = dockerfile.read_text() + self.assertIn("COPY mypkg ./mypkg", text) + self.assertIn("COPY runner.py ./runner.py", text) + + def test_add_copy_statements_raises_without_env(self): + dockerfile = self.project_path / "Dockerfile" + dockerfile.write_text("ENTRYPOINT test.sh\n") + with self.assertRaises(ValueError): + DockerHelper._add_copy_statements_to_dockerfile(str(dockerfile), []) + + def test_add_copy_statements_raises_without_entrypoint(self): + dockerfile = self.project_path / "Dockerfile" + dockerfile.write_text("ENV TEST=1\n") + with self.assertRaises(ValueError): + DockerHelper._add_copy_statements_to_dockerfile(str(dockerfile), []) + + def test_parse_minikube_env(self): + output = 'export FOO="bar"\nexport BAZ="qux"\n' + env = DockerHelper._parse_minikube_env(output) + self.assertEqual(env["FOO"], "bar") + self.assertEqual(env["BAZ"], "qux") + + @patch("gaiaflow.managers.minikube_manager.run") + def test_build_local(self, mock_run): + dockerfile = self._write_dockerfile() + self.helper._build_local(dockerfile) + mock_run.assert_called() + + @patch("gaiaflow.managers.minikube_manager.run") + @patch.object(MinikubeHelper, "run_cmd") + def test_build_minikube(self, mock_run_cmd, mock_run): + mock_run_cmd.return_value = MagicMock(stdout=b'export DOCKER_TLS_VERIFY="1"\n') + dockerfile = self._write_dockerfile() + self.helper._build_minikube(dockerfile) + mock_run.assert_called() + + @patch("gaiaflow.managers.minikube_manager.log_error") + def test_build_image_missing_dockerfile(self, mock_log_error): + bad_path = self.project_path / "Dockerfile" + self.assertFalse(bad_path.exists()) + with ( + patch.object(self.helper, "_update_dockerfile") as mock_update, + patch.object(self.helper, "_build_local") as mock_local, + patch.object(self.helper, "_build_minikube") as mock_minikube, + ): + self.helper.build_image(bad_path) + + mock_log_error.assert_called_once() + mock_update.assert_not_called() + mock_local.assert_not_called() + mock_minikube.assert_not_called() + + @patch( + "gaiaflow.managers.minikube_manager.find_python_packages", + return_value=["mypkg"], + ) + @patch("gaiaflow.managers.minikube_manager.temporary_copy") + def test_build_image_local(self, mock_temp_copy, _): + dockerfile = self._write_dockerfile() + self.helper.local = True + + with ( + patch.object(self.helper, "_update_dockerfile") as mock_update, + patch.object(self.helper, "_build_local") as mock_local, + ): + self.helper.build_image(dockerfile) + + mock_update.assert_called_once_with(dockerfile) + mock_local.assert_called_once_with(dockerfile) + mock_temp_copy.assert_called_once() # runner.py should be copied + + @patch( + "gaiaflow.managers.minikube_manager.find_python_packages", + return_value=["mypkg"], + ) + @patch("gaiaflow.managers.minikube_manager.temporary_copy") + def test_build_image_minikube(self, mock_temp_copy, _): + dockerfile = self._write_dockerfile() + self.helper.local = False + + with ( + patch.object(self.helper, "_update_dockerfile") as mock_update, + patch.object(self.helper, "_build_minikube") as mock_minikube, + ): + self.helper.build_image(dockerfile) + + mock_update.assert_called_once_with(dockerfile) + mock_minikube.assert_called_once_with(dockerfile) + mock_temp_copy.assert_called_once() + + +class TestKubeConfigHelper(unittest.TestCase): + def setUp(self): + self.tmpdir = tempfile.TemporaryDirectory() + self.gaia_path = Path(self.tmpdir.name) + (self.gaia_path / "docker_stuff").mkdir() + self.helper = KubeConfigHelper(gaiaflow_path=self.gaia_path, os_type="linux") + + def tearDown(self): + self.tmpdir.cleanup() + + def _write_kube_config(self, data): + kube_dir = Path.home() / ".kube" + kube_dir.mkdir(exist_ok=True) + kube_config = kube_dir / "config" + with open(kube_config, "w") as f: + yaml.dump(data, f) + return kube_config + + @patch("subprocess.call", return_value=0) + def test_write_inline_creates_file(self, _): + kube_config = self._write_kube_config({"clusters": []}) + self.helper._write_inline(kube_config) + out_file = self.gaia_path / "docker_stuff" / "kube_config_inline" + self.assertTrue(out_file.exists()) + + def test_backup_and_patch_config(self): + kube_config = self._write_kube_config( + {"clusters": [{"cluster": {"server": "127.0.0.1"}}]} + ) + backup = kube_config.with_suffix(".backup") + self.helper._backup_kube_config(kube_config, backup) + self.assertTrue(backup.exists()) + + self.helper._patch_kube_config(kube_config) + patched = yaml.safe_load(open(kube_config)) + self.assertIn("clusters", patched) + + @patch("subprocess.call", return_value=0) + @patch("gaiaflow.managers.minikube_manager.is_wsl", return_value=False) + def test_create_inline_linux(self, *_): + helper = KubeConfigHelper(self.gaia_path, os_type="linux") + self._write_kube_config({"clusters": [{"cluster": {"server": "127.0.0.1"}}]}) + helper.create_inline() + self.assertTrue( + (self.gaia_path / "docker_stuff" / "kube_config_inline").exists() + ) + + @patch("gaiaflow.managers.minikube_manager.is_wsl", return_value=False) + def test_create_inline_windows_branch(self, _): + helper = KubeConfigHelper(self.gaia_path, os_type="windows") + kube_config = self._write_kube_config({"clusters": []}) + backup_config = kube_config.with_suffix(".backup") + backup_config.write_text("backup") + + with ( + patch("shutil.copy") as mock_copy, + patch.object(Path, "unlink") as mock_unlink, + ): + helper.create_inline() + + mock_copy.assert_called_once_with(backup_config, kube_config) + mock_unlink.assert_called_once() + + @patch("gaiaflow.managers.minikube_manager.is_wsl", return_value=True) + def test_create_inline_wsl_branch(self, _): + helper = KubeConfigHelper(self.gaia_path, os_type="linux") + kube_config = self._write_kube_config({"clusters": []}) + backup_config = kube_config.with_suffix(".backup") + backup_config.write_text("backup") + + with ( + patch("shutil.copy") as mock_copy, + patch.object(Path, "unlink") as mock_unlink, + ): + helper.create_inline() + + mock_copy.assert_called_once_with(backup_config, kube_config) + mock_unlink.assert_called_once() + + @patch("gaiaflow.managers.minikube_manager.is_wsl", return_value=True) + def test_patch_wsl(self, _): + helper = KubeConfigHelper(self.gaia_path, os_type="linux") + config = self._write_kube_config({"clusters": [{"cluster": {"server": "localhost"}}]}) + helper._patch_kube_config(config) + data = yaml.safe_load(open(config)) + assert data["clusters"][0]["cluster"]["insecure-skip-tls-verify"] + + def test_patch_windows(self): + helper = KubeConfigHelper(self.gaia_path, os_type="windows") + config = self._write_kube_config({"clusters": [{"cluster": {"server": "127.0.0.1"}}]}) + helper._patch_kube_config(config) + data = yaml.safe_load(open(config)) + assert data["clusters"][0]["cluster"]["server"] == "host.docker.internal" + + +class TestMinikubeManager(unittest.TestCase): + def setUp(self): + self.tmpdir = tempfile.TemporaryDirectory() + self.manager = MinikubeManager( + gaiaflow_path=Path(self.tmpdir.name), + user_project_path=Path(self.tmpdir.name), + action=BaseAction.START, + local=True, + image_name="img", + ) + + + def tearDown(self): + self.tmpdir.cleanup() + + + def test_valid_actions(self): + actions = self.manager._get_valid_actions() + self.assertIn(ExtendedAction.DOCKERIZE, actions) + + + def test_run_raises_on_missing_action(self): + with self.assertRaises(ValueError): + MinikubeManager.run() + + + def test_run_raises_on_unknown_action(self): + with self.assertRaises(ValueError): + MinikubeManager.run(gaiaflow_path=Path(self.tmpdir.name), + user_project_path=Path(self.tmpdir.name),action="not-real") + + + @patch.object(MinikubeManager, "start") + def test_run_start_action(self, mock_start): + MinikubeManager.run( + gaiaflow_path=Path(self.tmpdir.name), + user_project_path=Path(self.tmpdir.name), + action=BaseAction.START, + ) + mock_start.assert_called() + + + @patch("subprocess.run", return_value=MagicMock(returncode=1)) + def test_create_secrets_creates_new(self, mock_run): + with patch("subprocess.check_call") as mock_call: + self.manager.create_secrets("new-secret", {"k": "v"}) + mock_call.assert_called() + + + @patch("subprocess.run", return_value=MagicMock(returncode=0)) + def test_create_secrets_skips_existing(self, _): + self.manager.create_secrets("existing-secret", {"k": "v"}) # should not raise + + + @patch("gaiaflow.managers.minikube_manager.run") + def test_cleanup_runs(self, mock_run): + self.manager.cleanup() + self.assertTrue(mock_run.called) + + @patch("gaiaflow.managers.minikube_manager.MlopsManager.run") + def test_stop_and_start_mlops(self, mock_run): + self.manager._stop_mlops() + self.manager._start_mlops() + assert mock_run.call_count == 2 + + @patch.object(KubeConfigHelper, "create_inline") + @patch.object(MinikubeHelper, "start") + @patch.object(MinikubeManager, "_stop_mlops") + @patch.object(MinikubeManager, "_start_mlops") + def test_start_calls_helpers( + self, mock_start_mlops, mock_stop_mlops, mock_mini_start, mock_inline + ): + self.manager.start() + mock_stop_mlops.assert_called() + mock_mini_start.assert_called() + mock_start_mlops.assert_called() + + @patch.object(MinikubeHelper, "stop") + def test_stop(self, mock_stop): + self.manager.stop() + mock_stop.assert_called() + + @patch.object(KubeConfigHelper, "create_inline") + def test_create_kube_config_inline(self, mock_inline): + self.manager.create_kube_config_inline() + mock_inline.assert_called() + + @patch.object(DockerHelper, "build_image") + def test_build_docker_image(self, mock_build): + self.manager.build_docker_image() + mock_build.assert_called() diff --git a/tests/managers/test_mlops_manager.py b/tests/managers/test_mlops_manager.py new file mode 100644 index 0000000..27c2da7 --- /dev/null +++ b/tests/managers/test_mlops_manager.py @@ -0,0 +1,350 @@ +import shutil +import socket +import tempfile +from pathlib import Path +from unittest import TestCase +from unittest.mock import patch, MagicMock + +import yaml + +from gaiaflow.constants import BaseAction, Service +from gaiaflow.managers.mlops_manager import ( + MlopsManager, + JupyterHelper, + DockerHelper, + DockerResources, +) + + +class TestMlopsManager(TestCase): + def setUp(self): + self.tmp_dir = tempfile.TemporaryDirectory() + self.base_path = Path(self.tmp_dir.name) + + self.user_project = self.base_path / "project" + self.user_project.mkdir() + (self.user_project / "environment.yml").write_text("name: test-env") + (self.user_project / "pyproject.toml").write_text("[project]\nname='test'") + (self.user_project / "dummy_package").mkdir() + (self.user_project / "dummy_package" / "__init__.py").write_text("") + + self.gaiaflow_context = self.base_path / "gaiaflow" + docker_dir = self.gaiaflow_context / "docker_stuff" / "docker-compose" + docker_dir.mkdir(parents=True) + (docker_dir / "docker-compose.yml").write_text( + yaml.dump({"x-airflow-common": {"volumes": ["./logs:/opt/airflow/logs"]}}) + ) + (docker_dir / "entrypoint.sh").write_text("#!/bin/bash\necho hi") + (self.gaiaflow_context / "docker_stuff" / "kube_config_inline").write_text("kube") + (self.gaiaflow_context / "environment.yml").write_text("name: test-env") + + self.manager = MlopsManager( + gaiaflow_path=self.gaiaflow_context, + user_project_path=self.user_project, + action=BaseAction.START, + service=Service.all, + ) + + def tearDown(self): + self.tmp_dir.cleanup() + + def test_invalid_env_tool(self): + with self.assertRaises(ValueError): + MlopsManager(self.gaiaflow_context, self.user_project, BaseAction.START, env_tool="pip") + + def test_unexpected_kwargs(self): + with self.assertRaises(TypeError): + MlopsManager(self.gaiaflow_context, self.user_project, BaseAction.START, bad_kwarg=True) + + @patch("gaiaflow.managers.mlops_manager.run") + @patch("gaiaflow.managers.mlops_manager.subprocess.Popen") + @patch("gaiaflow.managers.mlops_manager.env_exists") + def test_run_dispatches_start(self, mock_env_exists, mock_popen, + mock_run): + mock_env_exists.return_value = True + MlopsManager.run( + gaiaflow_path=self.manager.gaiaflow_path, + user_project_path=self.manager.user_project_path, + action=BaseAction.START, + service=Service.all, + ) + args, kwargs = mock_run.call_args + + self.assertIn("docker", args[0]) + self.assertIn("compose", args[0]) + self.assertIn("up", args[0]) + self.assertIn("-d", args[0]) + self.assertTrue(any("docker-compose.yml" in x for x in args[0])) + self.assertIn("Error running docker compose", args[1]) + + args, kwargs = mock_popen.call_args + print("subprocess.Popen args:", args) + print("subprocess.Popen kwargs:", kwargs) + + self.assertIn("jupyter", args[0]) + self.assertIn("lab", args[0]) + self.assertIn("test-env", args[0]) + self.assertIn("--port=8895", args[0]) + self.assertIn("mamba", args[0]) + + @patch("gaiaflow.managers.mlops_manager.run") + @patch("gaiaflow.managers.mlops_manager.subprocess.Popen") + @patch("gaiaflow.managers.mlops_manager.env_exists") + def test_run_dispatches_start_jupyter_custom_values(self, mock_env_exists, + mock_popen, + mock_run): + mock_env_exists.return_value = True + MlopsManager.run( + gaiaflow_path=self.manager.gaiaflow_path, + user_project_path=self.manager.user_project_path, + action=BaseAction.START, + service=Service.all, + jupyter_port=8181, + env_tool="conda", + ) + args, kwargs = mock_run.call_args + + self.assertIn("docker", args[0]) + self.assertIn("compose", args[0]) + self.assertIn("up", args[0]) + self.assertIn("-d", args[0]) + self.assertTrue(any("docker-compose.yml" in x for x in args[0])) + self.assertIn("Error running docker compose", args[1]) + + args, kwargs = mock_popen.call_args + print("subprocess.Popen args:", args) + print("subprocess.Popen kwargs:", kwargs) + + self.assertIn("jupyter", args[0]) + self.assertIn("lab", args[0]) + self.assertIn("test-env", args[0]) + self.assertIn("--port=8181", args[0]) + self.assertIn("conda", args[0]) + + def test_run_invalid_action(self): + with self.assertRaises(ValueError): + MlopsManager.run( + gaiaflow_path=self.manager.gaiaflow_path, + user_project_path=self.manager.user_project_path, + action="not-an-action", + ) + + @patch("gaiaflow.managers.mlops_manager.env_exists") + @patch("gaiaflow.managers.mlops_manager.subprocess.Popen") + def test_start_force_new(self, mock_popen, mock_env_exists): + self.manager.force_new = True + with patch.object(self.manager, "cleanup") as mock_cleanup: + self.manager.start() + mock_cleanup.assert_called_once() + + @patch("gaiaflow.managers.mlops_manager.env_exists") + @patch("gaiaflow.managers.mlops_manager.subprocess.Popen") + def test_start_service_jupyter(self, mock_popen, mock_env_exists): + self.manager.service = Service.jupyter + with patch.object(self.manager.jupyter, "check_port") as mock_check, \ + patch.object(self.manager.jupyter, "start") as mock_start: + self.manager.start() + mock_check.assert_called_once() + mock_start.assert_called_once() + + def test_start_single_service(self): + self.manager.service = Service.airflow + with patch.object(self.manager.docker, "run_compose") as mock_compose: + self.manager.start() + mock_compose.assert_called_with(["up", "-d"], Service.airflow) + + self.manager.service = Service.mlflow + with patch.object(self.manager.docker, "run_compose") as mock_compose: + self.manager.start() + mock_compose.assert_called_with(["up", "-d"], Service.mlflow) + + def test_run_missing_action_raises(self): + with self.assertRaises(ValueError) as ctx: + MlopsManager.run( + gaiaflow_path=self.manager.gaiaflow_path, + user_project_path=self.manager.user_project_path, + service=Service.all, + ) + self.assertIn("Missing required argument 'action'", str(ctx.exception)) + + @patch("gaiaflow.managers.mlops_manager.run") + @patch.object(MlopsManager, "_build_docker_images") + def test_start_triggers_build_docker_images(self, mock_build, mock_run): + self.manager.docker_build = True + self.manager.service = Service.airflow + self.manager.start() + mock_build.assert_called_once() + + def test_build_docker_images_cache_and_no_cache(self): + mock_docker = MagicMock() + self.manager.docker = mock_docker + mock_docker.run_compose = MagicMock() + + self.manager.cache = False + self.manager.service = Service.all + + self.manager._build_docker_images() + args, _ = mock_docker.run_compose.call_args + self.assertIn("--no-cache", args[0]) + + self.manager.cache = True + self.manager.service = Service.airflow + + self.manager._build_docker_images() + args, _ = mock_docker.run_compose.call_args + self.assertNotIn("--no-cache", args[0]) + + def test_stop_all(self): + with patch.object(self.manager, "jupyter") as mock_jupyter, \ + patch.object(self.manager, "docker") as mock_docker: + self.manager.stop() + mock_jupyter.stop.assert_called() + mock_docker.run_compose.assert_called() + + self.manager.delete_volume = True + self.manager.stop() + mock_jupyter.stop.assert_called() + mock_docker.run_compose.assert_called() + + def test_stop_jupyter(self): + self.manager.service = Service.jupyter + with patch.object(self.manager.jupyter, "stop") as mock_stop: + self.manager.stop() + mock_stop.assert_called_once() + + def test_stop_single_service_with_volume(self): + self.manager.service = Service.mlflow + self.manager.delete_volume = True + with patch.object(self.manager.docker, "run_compose") as mock_compose: + self.manager.stop() + args, _ = mock_compose.call_args + self.assertIn("-v", args[0]) + + def test_cleanup_with_prune(self): + self.manager.prune = True + with patch("shutil.rmtree") as mock_rm, \ + patch("gaiaflow.managers.mlops_manager.delete_project_state") as mock_del, \ + patch.object(self.manager.docker, "prune") as mock_prune: + self.manager.cleanup() + mock_rm.assert_called_once() + mock_del.assert_called_once() + mock_prune.assert_called_once() + + def test_cleanup_missing_context(self): + shutil.rmtree(self.gaiaflow_context) + with patch("gaiaflow.managers.mlops_manager.log_error") as mock_log: + self.manager.cleanup() + mock_log.assert_called() + + def test_delete_volume_logging(self): + self.manager.delete_volume = True + down_cmd = ["down"] + if self.manager.delete_volume: + down_cmd.append("-v") + self.assertIn("-v", down_cmd) + + def test_update_env_file_sets_uid(self): + env_path = self.base_path / ".env" + self.manager._update_env_file_with_airflow_uid(env_path) + self.assertIn("AIRFLOW_UID", env_path.read_text()) + + env_path.write_text("SOME_VAR=1\nAIRFLOW_UID=2\n") + self.manager._update_env_file_with_airflow_uid(env_path) + self.assertIn("SOME_VAR", env_path.read_text()) + self.assertIn("AIRFLOW_UID", env_path.read_text()) + + def test_update_env_file_updates_existing(self): + env_path = self.base_path / ".env" + env_path.write_text("AIRFLOW_UID=9999\n") + self.manager._update_env_file_with_airflow_uid(env_path) + content = env_path.read_text() + self.assertNotIn("9999", content) + + def test_update_files_rewrites_compose(self): + with patch("gaiaflow.managers.mlops_manager.find_python_packages", return_value=["dummy_package"]), \ + patch("gaiaflow.managers.mlops_manager.set_permissions"): + self.manager._update_files() + compose_path = self.gaiaflow_context / "docker_stuff" / "docker-compose" / "docker-compose.yml" + data = yaml.safe_load(compose_path.read_text()) + vols = data["x-airflow-common"]["volumes"] + self.assertTrue(any("dummy_package" in v for v in vols)) + self.assertTrue(any("/var/run/docker.sock" in v for v in vols)) + + @patch("gaiaflow.managers.mlops_manager.update_micromamba_env_in_docker") + def test_update_deps_calls_update_and_logs(self, mock_update): + MlopsManager.update_deps() + mock_update.assert_called_once_with(DockerResources.AIRFLOW_CONTAINERS) + + def test_jupyter_port_in_use(self): + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.bind(("127.0.0.1", 0)) + port = sock.getsockname()[1] + sock.listen(1) + + helper = JupyterHelper(port, "mamba", None, self.manager.gaiaflow_path) + with self.assertRaises(SystemExit): + helper.check_port() + sock.close() + + def test_jupyter_get_env_name(self): + helper = JupyterHelper(8895, "mamba", None, self.manager.gaiaflow_path) + name = helper.get_env_name() + self.assertEqual(name, "test-env") + + def test_jupyter_start_runs_subprocess(self): + helper = JupyterHelper(8895, "mamba", "custom-env", self.manager.gaiaflow_path) + with patch("subprocess.Popen") as mock_popen, \ + patch("gaiaflow.managers.mlops_manager.env_exists", return_value=True): + helper.start() + mock_popen.assert_called() + + @patch("psutil.process_iter") + def test_stop_jupyter_processes(self, mock_iter): + proc_mock = MagicMock() + proc_mock.info = {"pid": 123, "name": "jupyter", "cmdline": []} + mock_iter.return_value = [proc_mock] + self.manager.stop() + proc_mock.terminate.assert_called_once() + proc_mock.wait.assert_called_once_with(timeout=5) + + @patch("gaiaflow.managers.mlops_manager.env_exists", return_value=False) + @patch("subprocess.Popen") + def test_start_jupyter_env_not_exists(self, mock_popen, mock_env): + env_name = "test-env" + self.manager.get_env_name = MagicMock(return_value=env_name) + self.manager.jupyter.start() + mock_popen.assert_not_called() + + def test_docker_helper_builds_command(self): + helper = DockerHelper(self.manager.gaiaflow_path, is_prod_local=False) + cmd = helper._base_cmd() + self.assertIn("docker", cmd) + self.assertIn("compose", cmd) + self.assertEqual(cmd.count("-f"), 1) + + def test_docker_helper_builds_command_prod_local(self): + helper = DockerHelper(self.manager.gaiaflow_path, is_prod_local=True) + cmd = helper._base_cmd() + self.assertIn("docker", cmd) + self.assertIn("compose", cmd) + self.assertEqual(cmd.count("-f"), 2) + + + def test_docker_services_for_known_and_unknown(self): + helper = DockerHelper(Path("/tmp"), False) + self.assertIn("mlflow", helper.docker_services_for("mlflow")) + self.assertEqual(helper.docker_services_for("unknown"), []) + + @patch("gaiaflow.managers.mlops_manager.handle_error") + @patch("gaiaflow.managers.mlops_manager.run") + def test_run_compose_with_service_and_unknown_service(self, mock_run, mock_handle): + with patch.object(self.manager.docker, "docker_services_for", + return_value=[]): + self.manager.docker.run_compose(["up"], service="unknown_service") + mock_handle.assert_called_once() + + def test_docker_prune(self): + helper = self.manager.docker + with patch("gaiaflow.managers.mlops_manager.run") as mock_run: + helper.prune() + self.assertGreaterEqual(mock_run.call_count, len(DockerResources.IMAGES)) diff --git a/tests/managers/test_utils.py b/tests/managers/test_utils.py new file mode 100644 index 0000000..f9a89c9 --- /dev/null +++ b/tests/managers/test_utils.py @@ -0,0 +1,357 @@ +import json +import tempfile +import unittest +from pathlib import Path +import docker +from unittest.mock import patch, mock_open, MagicMock + +from gaiaflow.managers import utils + + +class TestUtils(unittest.TestCase): + def setUp(self): + self.tmp_dir = tempfile.TemporaryDirectory() + self.base_path = Path(self.tmp_dir.name) + self.gaiaflow_path = self.base_path / "gaiaflow-project" + self.gaiaflow_path.mkdir() + + required_structure = { + "docker_stuff": { + "docker-compose": [ + "docker-compose.yml", + "docker-compose-minikube-network.yml", + "entrypoint.sh", + ], + "airflow": ["Dockerfile"], + "mlflow": ["Dockerfile", "requirements.txt"], + "user-package": ["Dockerfile"], + "_files_": ["kube_config_inline"], + } + } + + self.gaiaflow_project_path = self.gaiaflow_path / "docker_stuff" + self.gaiaflow_project_path.mkdir(exist_ok=True) + for folder, contents in required_structure["docker_stuff"].items(): + if folder != "_files_": + folder_path = self.gaiaflow_project_path / folder + folder_path.mkdir(parents=True, exist_ok=True) + for file in contents: + (folder_path / file).write_text(f"dummy {file}") + else: + for file in contents: + (self.gaiaflow_project_path / file).write_text(f"dummy {file}") + + self.state_file_folder = self.base_path / ".gaiaflow" + self.state_file_folder.mkdir(exist_ok=True) + self.state_file = self.state_file_folder / "state.json" + self.state_data = {str(self.gaiaflow_path): {"project_path": str(self.base_path)}} + self.state_file.write_text(json.dumps(self.state_data)) + patcher = patch("gaiaflow.managers.utils.GAIAFLOW_STATE_FILE", + self.state_file) + self.addCleanup(patcher.stop) + patcher.start() + + def tearDown(self): + self.tmp_dir.cleanup() + + def test_get_gaialfow_version_from_pyproject(self): + with patch("importlib.metadata.version", side_effect=Exception): + pyproject_path = self.base_path / "pyproject.toml" + pyproject_path.write_text("[project]\nversion = '1.2.3'\n") + + with patch( + "pathlib.Path.read_text", return_value=pyproject_path.read_text() + ): + version = utils.get_gaiaflow_version() + self.assertEqual(version, "1.2.3") + + def test_path_exists_in_state_and_fs(self): + exists = utils.gaiaflow_path_exists_in_state(self.gaiaflow_path, check_fs=True) + self.assertTrue(exists) + + def test_path_missing_in_state(self): + missing_path = self.base_path / "missing" + exists = utils.gaiaflow_path_exists_in_state(missing_path) + self.assertFalse(exists) + + def test_path_exists_only_in_state(self): + missing_path = self.base_path / "nonexistent" + state = json.loads(self.state_file.read_text()) + state[str(missing_path)] = {"project_path": str(self.base_path)} + self.state_file.write_text(json.dumps(state)) + exists = utils.gaiaflow_path_exists_in_state(missing_path, check_fs=False) + self.assertTrue(exists) + + def test_check_fs_path_missing_on_disk(self): + missing_on_disk = self.base_path / "nonexistent" + state = json.loads(self.state_file.read_text()) + state[str(missing_on_disk)] = {"project_path": str(self.base_path)} + self.state_file.write_text(json.dumps(state)) + + with patch("typer.echo") as mock_echo: + exists = utils.gaiaflow_path_exists_in_state(missing_on_disk, check_fs=True) + self.assertFalse(exists) + mock_echo.assert_called_once() + self.assertIn( + "Gaiaflow path exists in state but not on disk", + mock_echo.call_args[0][0], + ) + + def test_check_fs_structure_invalid(self): + invalid_path = self.base_path / "bad_project" + invalid_path.mkdir() + state = json.loads(self.state_file.read_text()) + state[str(invalid_path)] = {"project_path": str(self.base_path)} + self.state_file.write_text(json.dumps(state)) + + with patch("gaiaflow.managers.utils.check_structure", + return_value=False): + exists = utils.gaiaflow_path_exists_in_state(invalid_path, check_fs=True) + self.assertFalse(exists) + + def test_invalid_state_file_returns_false(self): + self.state_file.write_text("not a json") + exists = utils.gaiaflow_path_exists_in_state(self.gaiaflow_path) + self.assertFalse(exists) + + def test_convert_crlf_to_lf(self): + with tempfile.NamedTemporaryFile(delete=False) as tmp: + tmp.write(b"line1\r\nline2\r\n") + tmp.close() + utils.convert_crlf_to_lf(tmp.name) + content = Path(tmp.name).read_bytes() + self.assertNotIn(b"\r\n", content) + self.assertIn(b"\n", content) + + def test_find_python_packages(self): + with tempfile.TemporaryDirectory() as tmpdir: + tmp = Path(tmpdir) + pkg = tmp / "pkg" + pkg.mkdir() + (pkg / "__init__.py").write_text("# package") + (tmp / "nopkg").mkdir() + packages = utils.find_python_packages(tmp) + self.assertIn("pkg", packages) + self.assertNotIn("nopkg", packages) + + def test_is_wsl_false(self): + with patch("builtins.open", side_effect=FileNotFoundError): + self.assertFalse(utils.is_wsl()) + + m = mock_open(read_data="Linux version ... Microsoft WSL2 ...") + with patch("builtins.open", m): + self.assertTrue(utils.is_wsl()) + + def test_log_info_and_error(self): + utils.log_info("info") + utils.log_error("error") + + @patch("subprocess.call", return_value=0) + def test_run_success(self, mock_call): + utils.run(["echo"], "fail") + mock_call.assert_called_once() + + @patch("gaiaflow.managers.utils.subprocess.call", + side_effect=FileNotFoundError("cmd not found")) + def test_run_fail(self, mock_call): + with self.assertRaises(FileNotFoundError): + utils.run(["nonexistent_command"], "fail") + mock_call.assert_called_once() + + def test_handle_error_exits(self): + with self.assertRaises(SystemExit): + utils.handle_error("fail") + + def test_save_and_load_project_state(self): + project_path = self.base_path / "proj" + project_path.mkdir() + gaiaflow_path = self.base_path / "gflow" + gaiaflow_path.mkdir() + utils.save_project_state(project_path, gaiaflow_path) + state = utils.load_project_state() + self.assertIn(str(gaiaflow_path), state) + + def test_save_with_no_project_state(self): + project_path = self.base_path / "proj" + project_path.mkdir() + self.state_file.unlink() + utils.save_project_state(project_path, self.gaiaflow_path) + + + def test_load_project_state_no_file(self): + if self.state_file.exists(): + self.state_file.unlink() + self.assertIsNone(utils.load_project_state()) + + def test_delete_project_state_removes_key(self): + project_path = self.base_path / "proj" + project_path.mkdir() + gaiaflow_path = self.base_path / "gflow" + gaiaflow_path.mkdir() + + utils.save_project_state(project_path, gaiaflow_path) + state = json.loads(self.state_file.read_text()) + self.assertIn(str(gaiaflow_path), state) + + utils.delete_project_state(gaiaflow_path) + state = json.loads(self.state_file.read_text()) + self.assertNotIn(str(gaiaflow_path), state) + + def test_state_file_missing(self): + self.state_file.unlink() + utils.delete_project_state(self.gaiaflow_path) + self.assertFalse(self.state_file.exists()) + + def test_delete_raises_jsondecodeerror(self): + self.state_file.write_text("{ invalid_data }") + with self.assertRaises(json.JSONDecodeError): + utils.delete_project_state(self.gaiaflow_path) + + def test_delete_raises_when_state_is_string(self): + self.state_file.write_text(json.dumps("invalid_data")) + with self.assertRaises(AttributeError): + utils.delete_project_state(self.gaiaflow_path) + + def test_update_project_state(self): + project_path = self.base_path / "proj" + project_path.mkdir() + gaiaflow_path = self.base_path / "gflow" + gaiaflow_path.mkdir() + + utils.save_project_state(project_path, gaiaflow_path) + state = json.loads(self.state_file.read_text()) + self.assertIn(str(gaiaflow_path), state) + + gaiaflow_path2 = self.base_path / "gflow_v2" + gaiaflow_path2.mkdir() + + utils.save_project_state(project_path, gaiaflow_path2) + state = json.loads(self.state_file.read_text()) + self.assertIn(str(gaiaflow_path2), state) + self.assertNotIn(str(gaiaflow_path), state) + + def test_check_structure_success_and_fail(self): + folder = self.base_path / "folder" + folder.mkdir() + (folder / "file.txt").write_text("hi") + structure = {"folder": ["file.txt"]} + self.assertTrue(utils.check_structure(self.base_path, structure)) + + structure = {"folder": ["missing.txt"]} + self.assertFalse(utils.check_structure(self.base_path, structure)) + + structure = {"missing_folder": ["file.txt"]} + self.assertFalse(utils.check_structure(self.base_path, structure)) + + nested_folder = folder / "nested" + nested_folder.mkdir() + (nested_folder / "nested_file.txt").write_text("ok") + structure = {"folder": {"nested": ["nested_file.txt"]}} + self.assertTrue(utils.check_structure(self.base_path, structure)) + + structure = {"folder": {"nested": ["missing_file.txt"]}} + self.assertFalse(utils.check_structure(self.base_path, structure)) + + (self.base_path / "file1.txt").write_text("ok") + structure = {"_files_": ["file1.txt"]} + self.assertTrue(utils.check_structure(self.base_path, structure)) + + structure = {"_files_": ["missing_file.txt"]} + self.assertFalse(utils.check_structure(self.base_path, structure)) + + def test_parse_key_value_pairs(self): + pairs = ["a=1", "b=2"] + result = utils.parse_key_value_pairs(pairs) + self.assertEqual(result, {"a": "1", "b": "2"}) + with self.assertRaises(Exception): + utils.parse_key_value_pairs(["invalid"]) + + def test_create_directory_and_permissions(self): + dir_path = self.base_path / "newdir" + utils.create_directory(str(dir_path)) + self.assertTrue(dir_path.exists()) + utils.set_permissions(dir_path) + + def test_create_directory_already_exists(self): + dir_path = self.base_path / "newdir" + utils.create_directory(str(dir_path)) + self.assertTrue(dir_path.exists()) + utils.create_directory(str(dir_path)) + + @patch("gaiaflow.managers.utils.fs") + def test_create_directory_makedirs_fails(self, mock_fs): + mock_fs.exists.return_value = False + mock_fs.makedirs.side_effect = Exception("boom") + with self.assertRaises(SystemExit): + utils.create_directory("fail_dir") + + @patch("subprocess.run") + def test_env_exists_true_false(self, mock_run): + mock_run.return_value.stdout = json.dumps({"envs": ["env1", "env2"]}) + self.assertTrue(utils.env_exists("env1")) + self.assertFalse(utils.env_exists("envX")) + + @patch("gaiaflow.managers.utils.docker.from_env") + @patch("gaiaflow.managers.utils.log_info") + @patch("gaiaflow.managers.utils.log_error") + def test_update_micromamba_env_in_docker(self, mock_error, mock_info, mock_docker): + client_mock = MagicMock() + mock_docker.return_value = client_mock + + container1 = MagicMock() + container1.exec_run.return_value = (0, b"success") + container2 = MagicMock() + container2.exec_run.return_value = (1, b"fail") + # container 3: container not found + # will raise docker.errors.NotFound + container4 = MagicMock() + container4.exec_run.side_effect = RuntimeError("boom") + + def get_container(name): + if name == "c1": + return container1 + if name == "c2": + return container2 + if name == "c3": + raise docker.errors.NotFound("not found") + if name == "c4": + return container4 + raise ValueError("Unexpected container") + + client_mock.containers.get.side_effect = get_container + + containers = ["c1", "c2", "c3", "c4"] + utils.update_micromamba_env_in_docker(containers, env_name="test_env", max_workers=4) + + client_mock.containers.get.assert_any_call("c1") + client_mock.containers.get.assert_any_call("c2") + client_mock.containers.get.assert_any_call("c3") + client_mock.containers.get.assert_any_call("c4") + + container1.exec_run.assert_called_once() + container2.exec_run.assert_called_once() + container4.exec_run.assert_called_once() + + mock_info.assert_any_call("[c1] Updated successfully.") + mock_error.assert_any_call("[c2] micromamba failed: fail") + mock_error.assert_any_call("Container 'c3' not found. Skipping.") + self.assertTrue( + any("[c4] Unexpected error: boom" in args[0] for args, + _ in mock_error.call_args_list) + ) + + def test_project_path_not_found(self): + missing_path = Path(self.tmp_dir.name) / "missing" + with self.assertRaises(FileNotFoundError) as ctx: + utils.create_gaiaflow_context_path(missing_path) + self.assertIn("not found", str(ctx.exception)) + + @patch("gaiaflow.managers.utils.get_gaiaflow_version", return_value="1.2.3") + def test_successful_context_path(self, mock_version): + project_path = self.base_path / "proj" + project_path.mkdir() + gaiaflow_path, user_project_path = utils.create_gaiaflow_context_path(project_path) + self.assertEqual(user_project_path, project_path.resolve()) + self.assertIn("gaiaflow-1.2.3-proj", str(gaiaflow_path)) + self.assertEqual(gaiaflow_path.parent, Path(tempfile.gettempdir())) + mock_version.assert_called_once() \ No newline at end of file diff --git a/tests/test_testing.py b/tests/test_testing.py new file mode 100644 index 0000000..f875a34 --- /dev/null +++ b/tests/test_testing.py @@ -0,0 +1,20 @@ +import os +import unittest + +from gaiaflow.testing import set_env_cm + + +class TestingTest(unittest.TestCase): + def test_set_env_cm(self): + old_env = dict(os.environ) + with set_env_cm(PARAM_example="abc", PARAM_EXAMPLE="xyz"): + self.assertEqual("abc", os.environ.get("PARAM_example")) + self.assertEqual("xyz", os.environ.get("PARAM_EXAMPLE")) + self.assertNotEqual(old_env, os.environ) + with set_env_cm(PARAM_example=None, PARAM_EXAMPLE=None): + self.assertEqual(None, os.environ.get("PARAM_example")) + self.assertEqual(None, os.environ.get("PARAM_EXAMPLE")) + self.assertEqual("abc", os.environ.get("PARAM_example")) + self.assertEqual("xyz", os.environ.get("PARAM_EXAMPLE")) + self.assertNotEqual(old_env, os.environ) + self.assertEqual(old_env, os.environ)