From 317962056110c237399bf8834f66f3db8533ac65 Mon Sep 17 00:00:00 2001 From: Yogesh Kumar Baljeet Singh Date: Mon, 15 Sep 2025 12:19:27 +0200 Subject: [PATCH 01/15] update CI --- .github/workflows/unittest.yml | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/.github/workflows/unittest.yml b/.github/workflows/unittest.yml index adf770d..747c90b 100644 --- a/.github/workflows/unittest.yml +++ b/.github/workflows/unittest.yml @@ -34,10 +34,11 @@ jobs: - name: Install and Test run: - pixi run test --cov=gaiaflow --cov-report=xml + pixi run test --cov=gaiaflow --cov-branch --cov-report=xml - name: Upload coverage reports to Codecov - uses: codecov/codecov-action@v4 + uses: codecov/codecov-action@v5 with: verbose: true - token: ${{ secrets.CODECOV_TOKEN }} \ No newline at end of file + token: ${{ secrets.CODECOV_TOKEN }} + slug: bcdev/gaiaflow \ No newline at end of file From bc484d4b0f32c55c13e28e194e2ef222f106cb38 Mon Sep 17 00:00:00 2001 From: Yogesh Kumar Baljeet Singh Date: Mon, 15 Sep 2025 12:42:06 +0200 Subject: [PATCH 02/15] renamed docker_stuff to _docker --- pixi.lock | 2 +- pyproject.toml | 4 ++-- src/{docker_stuff => _docker}/airflow/Dockerfile | 0 .../docker-compose-minikube-network.yml | 0 .../docker-compose/docker-compose.yml | 4 ++-- .../docker-compose/entrypoint.sh | 0 src/{docker_stuff => _docker}/kube_config_inline | 0 src/{docker_stuff => _docker}/mlflow/Dockerfile | 0 .../mlflow/requirements.txt | 0 .../user-package/Dockerfile | 0 src/gaiaflow/managers/minikube_manager.py | 6 +++--- src/gaiaflow/managers/mlops_manager.py | 16 ++++++++-------- src/gaiaflow/managers/utils.py | 2 +- tests/managers/test_minikube_manager.py | 6 +++--- tests/managers/test_mlops_manager.py | 6 +++--- tests/managers/test_utils.py | 6 +++--- 16 files changed, 26 insertions(+), 26 deletions(-) rename src/{docker_stuff => _docker}/airflow/Dockerfile (100%) rename src/{docker_stuff => _docker}/docker-compose/docker-compose-minikube-network.yml (100%) rename src/{docker_stuff => _docker}/docker-compose/docker-compose.yml (99%) rename src/{docker_stuff => _docker}/docker-compose/entrypoint.sh (100%) rename src/{docker_stuff => _docker}/kube_config_inline (100%) rename src/{docker_stuff => _docker}/mlflow/Dockerfile (100%) rename src/{docker_stuff => _docker}/mlflow/requirements.txt (100%) rename src/{docker_stuff => _docker}/user-package/Dockerfile (100%) diff --git a/pixi.lock b/pixi.lock index fa3db2e..aac4e30 100644 --- a/pixi.lock +++ b/pixi.lock @@ -2383,7 +2383,7 @@ packages: - pypi: ./ name: gaiaflow version: 0.0.4.dev0 - sha256: 5e3ceaa3916b7d3961eb01d16f9b1ee62a9c146a0ae379f3aab37f3fe54b1b78 + sha256: ccd33e40228db9872b498169a03a16977eeb05cb1e99d6f744f061e879c36529 requires_dist: - typer>=0.16.0,<0.17 - fsspec>=2025.7.0,<2026 diff --git a/pyproject.toml b/pyproject.toml index 4060124..1c35aa6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,12 +31,12 @@ build-backend = "hatchling.build" requires = ["hatchling"] [tool.hatch.build.targets.wheel] -packages = ["src/gaiaflow", "src/docker_stuff"] +packages = ["src/gaiaflow", "src/_docker"] [tool.hatch.build] include = [ "src/gaiaflow", - "src/docker_stuff", + "src/_docker", ] [tool.mypy] diff --git a/src/docker_stuff/airflow/Dockerfile b/src/_docker/airflow/Dockerfile similarity index 100% rename from src/docker_stuff/airflow/Dockerfile rename to src/_docker/airflow/Dockerfile diff --git a/src/docker_stuff/docker-compose/docker-compose-minikube-network.yml b/src/_docker/docker-compose/docker-compose-minikube-network.yml similarity index 100% rename from src/docker_stuff/docker-compose/docker-compose-minikube-network.yml rename to src/_docker/docker-compose/docker-compose-minikube-network.yml diff --git a/src/docker_stuff/docker-compose/docker-compose.yml b/src/_docker/docker-compose/docker-compose.yml similarity index 99% rename from src/docker_stuff/docker-compose/docker-compose.yml rename to src/_docker/docker-compose/docker-compose.yml index d8258ad..d4b203d 100644 --- a/src/docker_stuff/docker-compose/docker-compose.yml +++ b/src/_docker/docker-compose/docker-compose.yml @@ -2,7 +2,7 @@ x-airflow-common: &airflow-common build: context: ../../ - dockerfile: "./docker_stuff/airflow/Dockerfile" + dockerfile: "./_docker/airflow/Dockerfile" environment: &airflow-common-env # NOTE: The following secret can be commited to .git as it is only for local development. @@ -75,7 +75,7 @@ services: mlflow: build: context: .. - dockerfile: "../docker_stuff/mlflow/Dockerfile" + dockerfile: "../_docker/mlflow/Dockerfile" container_name: mlflow ports: - 5000:5000 diff --git a/src/docker_stuff/docker-compose/entrypoint.sh b/src/_docker/docker-compose/entrypoint.sh similarity index 100% rename from src/docker_stuff/docker-compose/entrypoint.sh rename to src/_docker/docker-compose/entrypoint.sh diff --git a/src/docker_stuff/kube_config_inline b/src/_docker/kube_config_inline similarity index 100% rename from src/docker_stuff/kube_config_inline rename to src/_docker/kube_config_inline diff --git a/src/docker_stuff/mlflow/Dockerfile b/src/_docker/mlflow/Dockerfile similarity index 100% rename from src/docker_stuff/mlflow/Dockerfile rename to src/_docker/mlflow/Dockerfile diff --git a/src/docker_stuff/mlflow/requirements.txt b/src/_docker/mlflow/requirements.txt similarity index 100% rename from src/docker_stuff/mlflow/requirements.txt rename to src/_docker/mlflow/requirements.txt diff --git a/src/docker_stuff/user-package/Dockerfile b/src/_docker/user-package/Dockerfile similarity index 100% rename from src/docker_stuff/user-package/Dockerfile rename to src/_docker/user-package/Dockerfile diff --git a/src/gaiaflow/managers/minikube_manager.py b/src/gaiaflow/managers/minikube_manager.py index 1182a36..247e3ef 100644 --- a/src/gaiaflow/managers/minikube_manager.py +++ b/src/gaiaflow/managers/minikube_manager.py @@ -268,7 +268,7 @@ def _patch_kube_config(self, kube_config: Path): yaml.dump(config_data, f) def _write_inline(self, kube_config: Path): - filename = self.gaiaflow_path / "docker_stuff" / "kube_config_inline" + filename = self.gaiaflow_path / "_docker" / "kube_config_inline" log_info("Creating kube config inline file...") with open(filename, "w") as f: subprocess.call( @@ -282,7 +282,7 @@ def _write_inline(self, kube_config: Path): "--minify", "--raw", ], - cwd=self.gaiaflow_path / "docker_stuff", + cwd=self.gaiaflow_path / "_docker", stdout=f, ) log_info(f"Created kube config inline file {filename}") @@ -393,7 +393,7 @@ def create_kube_config_inline(self): def build_docker_image(self): dockerfile_path = ( - self.gaiaflow_path / "docker_stuff" / "user-package" / "Dockerfile" + self.gaiaflow_path / "_docker" / "user-package" / "Dockerfile" ) self.docker_helper.build_image(dockerfile_path) diff --git a/src/gaiaflow/managers/mlops_manager.py b/src/gaiaflow/managers/mlops_manager.py index 6929c22..3af68c8 100644 --- a/src/gaiaflow/managers/mlops_manager.py +++ b/src/gaiaflow/managers/mlops_manager.py @@ -82,12 +82,12 @@ def _base_cmd(self) -> list[str]: "docker", "compose", "-f", - f"{self.gaiaflow_path}/docker_stuff/docker-compose/docker-compose.yml", + f"{self.gaiaflow_path}/_docker/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", + f"{self.gaiaflow_path}/_docker/docker-compose/docker-compose-minikube-network.yml", ] return base @@ -403,10 +403,10 @@ def _create_gaiaflow_context(self): self.fs.makedirs(self.gaiaflow_path, exist_ok=True) package_dir = Path(__file__).parent.parent.resolve() - docker_dir = package_dir.parent / "docker_stuff" + docker_dir = package_dir.parent / "_docker" shutil.copytree( - docker_dir, self.gaiaflow_path / "docker_stuff", dirs_exist_ok=True + docker_dir, self.gaiaflow_path / "_docker", dirs_exist_ok=True ) log_info(f"Gaiaflow context created at {self.gaiaflow_path}") @@ -445,11 +445,11 @@ def _collect_volumes(self, compose_data: dict) -> list[str]: # Add special mounts for prod_local mode kube_config = ( - self.gaiaflow_path.resolve() / "docker_stuff" / "kube_config_inline" + self.gaiaflow_path.resolve() / "_docker" / "kube_config_inline" ).as_posix() entrypoint = ( self.gaiaflow_path.resolve() - / "docker_stuff" + / "_docker" / "docker-compose" / "entrypoint.sh" ).as_posix() @@ -472,7 +472,7 @@ def _update_files(self): compose_path = ( self.gaiaflow_path - / "docker_stuff" + / "_docker" / "docker-compose" / "docker-compose.yml" ) @@ -487,7 +487,7 @@ def _update_files(self): yaml.dump(compose_data, f) entrypoint_path = ( - self.gaiaflow_path / "docker_stuff" / "docker-compose" / "entrypoint.sh" + self.gaiaflow_path / "_docker" / "docker-compose" / "entrypoint.sh" ) set_permissions(entrypoint_path) convert_crlf_to_lf(entrypoint_path) diff --git a/src/gaiaflow/managers/utils.py b/src/gaiaflow/managers/utils.py index 53a0f52..7eb89b8 100644 --- a/src/gaiaflow/managers/utils.py +++ b/src/gaiaflow/managers/utils.py @@ -140,7 +140,7 @@ def check_structure(base_path: Path, structure: dict) -> bool: def gaiaflow_path_exists_in_state(gaiaflow_path: Path, check_fs: bool = True) -> bool: REQUIRED_STRUCTURE = { - "docker_stuff": { + "_docker": { "docker-compose": [ "docker-compose.yml", "docker-compose-minikube-network.yml", diff --git a/tests/managers/test_minikube_manager.py b/tests/managers/test_minikube_manager.py index aeb6723..05f39f5 100644 --- a/tests/managers/test_minikube_manager.py +++ b/tests/managers/test_minikube_manager.py @@ -251,7 +251,7 @@ 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.gaia_path / "_docker").mkdir() self.helper = KubeConfigHelper(gaiaflow_path=self.gaia_path, os_type="linux") def tearDown(self): @@ -269,7 +269,7 @@ def _write_kube_config(self, data): 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" + out_file = self.gaia_path / "_docker" / "kube_config_inline" self.assertTrue(out_file.exists()) def test_backup_and_patch_config(self): @@ -291,7 +291,7 @@ def test_create_inline_linux(self, *_): 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() + (self.gaia_path / "_docker" / "kube_config_inline").exists() ) @patch("gaiaflow.managers.minikube_manager.is_wsl", return_value=False) diff --git a/tests/managers/test_mlops_manager.py b/tests/managers/test_mlops_manager.py index 27c2da7..bf560a2 100644 --- a/tests/managers/test_mlops_manager.py +++ b/tests/managers/test_mlops_manager.py @@ -29,13 +29,13 @@ def setUp(self): (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 = self.gaiaflow_context / "_docker" / "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 / "_docker" / "kube_config_inline").write_text("kube") (self.gaiaflow_context / "environment.yml").write_text("name: test-env") self.manager = MlopsManager( @@ -264,7 +264,7 @@ 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" + compose_path = self.gaiaflow_context / "_docker" / "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)) diff --git a/tests/managers/test_utils.py b/tests/managers/test_utils.py index f9a89c9..7a586ad 100644 --- a/tests/managers/test_utils.py +++ b/tests/managers/test_utils.py @@ -16,7 +16,7 @@ def setUp(self): self.gaiaflow_path.mkdir() required_structure = { - "docker_stuff": { + "_docker": { "docker-compose": [ "docker-compose.yml", "docker-compose-minikube-network.yml", @@ -29,9 +29,9 @@ def setUp(self): } } - self.gaiaflow_project_path = self.gaiaflow_path / "docker_stuff" + self.gaiaflow_project_path = self.gaiaflow_path / "_docker" self.gaiaflow_project_path.mkdir(exist_ok=True) - for folder, contents in required_structure["docker_stuff"].items(): + for folder, contents in required_structure["_docker"].items(): if folder != "_files_": folder_path = self.gaiaflow_project_path / folder folder_path.mkdir(parents=True, exist_ok=True) From e08d76ada11fe8d8a2b77ad79781de8a2a76c7f7 Mon Sep 17 00:00:00 2001 From: Yogesh Kumar Baljeet Singh Date: Mon, 15 Sep 2025 12:42:18 +0200 Subject: [PATCH 03/15] add codecov to README --- README.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/README.md b/README.md index e64ce22..a92588b 100644 --- a/README.md +++ b/README.md @@ -1,9 +1,11 @@ # Gaiaflow ![PyPI - Version](https://img.shields.io/pypi/v/gaiaflow) +[![codecov](https://codecov.io/gh/bcdev/gaiaflow/graph/badge.svg?token=pc9DfJx6bu)](https://codecov.io/gh/bcdev/gaiaflow) [![Pixi Badge](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/prefix-dev/pixi/main/assets/badge/v0.json)](https://pixi.sh) [![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/charliermarsh/ruff/main/assets/badge/v0.json)](https://github.com/charliermarsh/ruff) [![Docs](https://img.shields.io/badge/docs-mkdocs-blue)](https://bcdev.github.io/gaiaflow/) + ![Static Badge](https://img.shields.io/badge/Airflow-3.0-8A2BE2?logo=apacheairflow) ![Static Badge](https://img.shields.io/badge/MLFlow-darkblue?logo=mlflow) ![Static Badge](https://img.shields.io/badge/MinIO-red?logo=minio) From 36fe1801d742dbaa80f504540608689860d6c77d Mon Sep 17 00:00:00 2001 From: Yogesh Kumar Baljeet Singh Date: Mon, 15 Sep 2025 12:57:45 +0200 Subject: [PATCH 04/15] rename DockerHelper to DockerComposeHelper in MlopsManager and ruff --- src/gaiaflow/managers/minikube_manager.py | 4 +--- src/gaiaflow/managers/mlops_manager.py | 13 ++++--------- tests/managers/test_mlops_manager.py | 8 ++++---- 3 files changed, 9 insertions(+), 16 deletions(-) diff --git a/src/gaiaflow/managers/minikube_manager.py b/src/gaiaflow/managers/minikube_manager.py index 247e3ef..4964627 100644 --- a/src/gaiaflow/managers/minikube_manager.py +++ b/src/gaiaflow/managers/minikube_manager.py @@ -392,9 +392,7 @@ def create_kube_config_inline(self): self.kube_helper.create_inline() def build_docker_image(self): - dockerfile_path = ( - self.gaiaflow_path / "_docker" / "user-package" / "Dockerfile" - ) + dockerfile_path = self.gaiaflow_path / "_docker" / "user-package" / "Dockerfile" self.docker_helper.build_image(dockerfile_path) def create_secrets(self, secret_name: str, secret_data: dict[str, Any]): diff --git a/src/gaiaflow/managers/mlops_manager.py b/src/gaiaflow/managers/mlops_manager.py index 3af68c8..b54b6c2 100644 --- a/src/gaiaflow/managers/mlops_manager.py +++ b/src/gaiaflow/managers/mlops_manager.py @@ -72,7 +72,7 @@ class DockerResources: } -class DockerHelper: +class DockerComposeHelper: def __init__(self, gaiaflow_path: Path, is_prod_local: bool): self.gaiaflow_path = gaiaflow_path self.is_prod_local = is_prod_local @@ -224,7 +224,7 @@ def __init__( self.user_env_name = user_env_name self.env_tool = env_tool - self.docker = DockerHelper(gaiaflow_path, prod_local) + self.docker = DockerComposeHelper(gaiaflow_path, prod_local) self.jupyter = JupyterHelper( jupyter_port, env_tool, user_env_name, gaiaflow_path ) @@ -405,9 +405,7 @@ def _create_gaiaflow_context(self): package_dir = Path(__file__).parent.parent.resolve() docker_dir = package_dir.parent / "_docker" - shutil.copytree( - docker_dir, self.gaiaflow_path / "_docker", dirs_exist_ok=True - ) + shutil.copytree(docker_dir, self.gaiaflow_path / "_docker", dirs_exist_ok=True) log_info(f"Gaiaflow context created at {self.gaiaflow_path}") def _collect_volumes(self, compose_data: dict) -> list[str]: @@ -471,10 +469,7 @@ def _update_files(self): yaml.preserve_quotes = True compose_path = ( - self.gaiaflow_path - / "_docker" - / "docker-compose" - / "docker-compose.yml" + self.gaiaflow_path / "_docker" / "docker-compose" / "docker-compose.yml" ) with open(compose_path) as f: diff --git a/tests/managers/test_mlops_manager.py b/tests/managers/test_mlops_manager.py index bf560a2..b1455bb 100644 --- a/tests/managers/test_mlops_manager.py +++ b/tests/managers/test_mlops_manager.py @@ -11,7 +11,7 @@ from gaiaflow.managers.mlops_manager import ( MlopsManager, JupyterHelper, - DockerHelper, + DockerComposeHelper, DockerResources, ) @@ -316,14 +316,14 @@ def test_start_jupyter_env_not_exists(self, mock_popen, mock_env): mock_popen.assert_not_called() def test_docker_helper_builds_command(self): - helper = DockerHelper(self.manager.gaiaflow_path, is_prod_local=False) + helper = DockerComposeHelper(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) + helper = DockerComposeHelper(self.manager.gaiaflow_path, is_prod_local=True) cmd = helper._base_cmd() self.assertIn("docker", cmd) self.assertIn("compose", cmd) @@ -331,7 +331,7 @@ def test_docker_helper_builds_command_prod_local(self): def test_docker_services_for_known_and_unknown(self): - helper = DockerHelper(Path("/tmp"), False) + helper = DockerComposeHelper(Path("/tmp"), False) self.assertIn("mlflow", helper.docker_services_for("mlflow")) self.assertEqual(helper.docker_services_for("unknown"), []) From acbd2ab9135588c20adcb35240241244cf24440d Mon Sep 17 00:00:00 2001 From: Yogesh Kumar Baljeet Singh Date: Mon, 15 Sep 2025 17:34:51 +0200 Subject: [PATCH 05/15] remove required project_path in cli --- src/gaiaflow/cli/commands/minikube.py | 20 ++++++++++---------- src/gaiaflow/cli/commands/mlops.py | 13 +++++++------ 2 files changed, 17 insertions(+), 16 deletions(-) diff --git a/src/gaiaflow/cli/commands/minikube.py b/src/gaiaflow/cli/commands/minikube.py index 755c41f..1979bb2 100644 --- a/src/gaiaflow/cli/commands/minikube.py +++ b/src/gaiaflow/cli/commands/minikube.py @@ -31,7 +31,6 @@ def load_imports(): @app.command(help="Start Gaiaflow production-like services.") def start( - project_path: Path = typer.Option(..., "--path", "-p", help="Path to your project"), force_new: bool = typer.Option( False, "--force-new", @@ -42,6 +41,7 @@ def start( ): """""" imports = load_imports() + project_path = Path.cwd() gaiaflow_path, user_project_path = imports.create_gaiaflow_context_path( project_path ) @@ -59,9 +59,9 @@ def start( @app.command(help="Stop Gaiaflow production-like services.") def stop( - project_path: Path = typer.Option(..., "--path", "-p", help="Path to your project"), ): imports = load_imports() + project_path = Path.cwd() gaiaflow_path, user_project_path = imports.create_gaiaflow_context_path( project_path ) @@ -78,7 +78,6 @@ def stop( @app.command(help="Restart Gaiaflow production-like services.") def restart( - project_path: Path = typer.Option(..., "--path", "-p", help="Path to your project"), force_new: bool = typer.Option( False, "--force-new", @@ -88,6 +87,7 @@ def restart( ), ): imports = load_imports() + project_path = Path.cwd() gaiaflow_path, user_project_path = imports.create_gaiaflow_context_path( project_path ) @@ -99,6 +99,7 @@ def restart( gaiaflow_path=gaiaflow_path, user_project_path=user_project_path, action=imports.BaseAction.RESTART, + force_new=force_new ) @@ -106,12 +107,12 @@ def restart( 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.") ), ): imports = load_imports() + project_path = Path.cwd() gaiaflow_path, user_project_path = imports.create_gaiaflow_context_path( project_path ) @@ -133,9 +134,9 @@ def dockerize( "cluster. To be used only when debugging required." ) def create_config( - project_path: Path = typer.Option(..., "--path", "-p", help="Path to your project"), ): imports = load_imports() + project_path = Path.cwd() gaiaflow_path, user_project_path = imports.create_gaiaflow_context_path( project_path ) @@ -152,13 +153,13 @@ def create_config( @app.command(help="Create secrets to provide to the production-like environment.") def create_secret( - project_path: Path = typer.Option(..., "--path", "-p", help="Path to your project"), name: str = typer.Option(..., "--name", help="Name of the secret"), data: list[str] = typer.Option( ..., "--data", help="Secret data as key=value pairs" ), ): imports = load_imports() + project_path = Path.cwd() secret_data = imports.parse_key_value_pairs(data) print(secret_data, name) gaiaflow_path, user_project_path = imports.create_gaiaflow_context_path( @@ -179,12 +180,11 @@ def create_secret( @app.command( help="Clean Gaiaflow production-like services. This will only remove the " - "minikube speicifc things. To remove local docker stuff, use the dev mode." + "minikube specific things. To remove local docker stuff, use the dev mode." ) -def cleanup( - project_path: Path = typer.Option(..., "--path", "-p", help="Path to your project"), -): +def cleanup(): imports = load_imports() + project_path = Path.cwd() gaiaflow_path, user_project_path = imports.create_gaiaflow_context_path( project_path ) diff --git a/src/gaiaflow/cli/commands/mlops.py b/src/gaiaflow/cli/commands/mlops.py index 5962a8b..3b68b78 100644 --- a/src/gaiaflow/cli/commands/mlops.py +++ b/src/gaiaflow/cli/commands/mlops.py @@ -34,7 +34,6 @@ def load_imports(): @app.command(help="Start Gaiaflow development services") def start( - project_path: Path = typer.Option(..., "--path", "-p", help="Path to your project"), force_new: bool = typer.Option( False, "--force-new", @@ -72,6 +71,7 @@ def start( ), ): imports = load_imports() + project_path = Path.cwd() typer.echo(f"Selected Gaiaflow services: {service}") gaiaflow_path, user_project_path = imports.create_gaiaflow_context_path( project_path @@ -118,7 +118,6 @@ def start( @app.command(help="Stop Gaiaflow development services") def stop( - project_path: Path = typer.Option(..., "--path", "-p", help="Path to your project"), service: List[Service] = typer.Option( ["all"], "--service", @@ -131,6 +130,7 @@ def stop( ): """""" imports = load_imports() + project_path = Path.cwd() gaiaflow_path, user_project_path = imports.create_gaiaflow_context_path( project_path ) @@ -160,7 +160,6 @@ def stop( @app.command(help="Restart Gaiaflow development services") def restart( - project_path: Path = typer.Option(..., "--path", "-p", help="Path to your project"), force_new: bool = typer.Option( False, "--force-new", @@ -189,6 +188,7 @@ def restart( ): """""" imports = load_imports() + project_path = Path.cwd() gaiaflow_path, user_project_path = imports.create_gaiaflow_context_path( project_path ) @@ -230,12 +230,12 @@ def restart( "also remove the state for this project." ) def cleanup( - project_path: Path = typer.Option(..., "--path", "-p", help="Path to your project"), prune: bool = typer.Option( False, "--prune", help="Prune Docker image, network and cache" ), ): imports = load_imports() + project_path = Path.cwd() gaiaflow_path, user_project_path = imports.create_gaiaflow_context_path( project_path ) @@ -252,15 +252,16 @@ 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.") ), ): imports = load_imports() + project_path = Path.cwd() gaiaflow_path, user_project_path = imports.create_gaiaflow_context_path( project_path ) + gaiaflow_path_exists = imports.gaiaflow_path_exists_in_state(gaiaflow_path, True) if not gaiaflow_path_exists: imports.save_project_state(user_project_path, gaiaflow_path) @@ -289,9 +290,9 @@ def dockerize( "its contents." ) def update_deps( - project_path: Path = typer.Option(..., "--path", "-p", help="Path to your project"), ): imports = load_imports() + project_path = Path.cwd() gaiaflow_path, user_project_path = imports.create_gaiaflow_context_path( project_path ) From 7c8de9de8d56193fca925e7460b694bcc38f2324 Mon Sep 17 00:00:00 2001 From: Yogesh Kumar Baljeet Singh Date: Mon, 15 Sep 2025 17:35:17 +0200 Subject: [PATCH 06/15] mount the user directory instead of figuring out the python paths --- src/gaiaflow/managers/mlops_manager.py | 22 ++++++---------------- src/gaiaflow/managers/utils.py | 22 +++++++++++++++++++++- 2 files changed, 27 insertions(+), 17 deletions(-) diff --git a/src/gaiaflow/managers/mlops_manager.py b/src/gaiaflow/managers/mlops_manager.py index b54b6c2..69198f7 100644 --- a/src/gaiaflow/managers/mlops_manager.py +++ b/src/gaiaflow/managers/mlops_manager.py @@ -37,6 +37,7 @@ save_project_state, set_permissions, update_micromamba_env_in_docker, + update_entrypoint_install_path, ) @@ -422,24 +423,11 @@ def _collect_volumes(self, compose_data: dict) -> list[str]: ) 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) - # 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 - ): - dst_path = f"/opt/airflow/{child.name}" - new_volumes.append(f"{child.resolve().as_posix()}:{dst_path}") + set_permissions(self.user_project_path, 0o755) + + new_volumes.append(f"{self.user_project_path.resolve().as_posix()}:/opt/airflow/{self.user_project_path.name}") # Add special mounts for prod_local mode kube_config = ( @@ -484,6 +472,8 @@ def _update_files(self): entrypoint_path = ( self.gaiaflow_path / "_docker" / "docker-compose" / "entrypoint.sh" ) + update_entrypoint_install_path(entrypoint_path, + str(self.user_project_path.name)) set_permissions(entrypoint_path) convert_crlf_to_lf(entrypoint_path) diff --git a/src/gaiaflow/managers/utils.py b/src/gaiaflow/managers/utils.py index 7eb89b8..3fcd0a7 100644 --- a/src/gaiaflow/managers/utils.py +++ b/src/gaiaflow/managers/utils.py @@ -225,7 +225,9 @@ def set_permissions(path, mode=0o777): fs.chmod(path, mode) log_info(f"Set permissions for {path}") except Exception: - log_info(f"Warning: Could not set permissions for {path}") + log_error(f"Warning: Could not set permissions for {path}") + log_error(f"Try running this command manually:\n chmod -R {mode:o} {path}") + log_info("Continuing...") def create_gaiaflow_context_path(project_path: Path) -> tuple[Path, Path]: @@ -304,3 +306,21 @@ def _update_one(cname: str): future.result() except Exception as e: log_error(f"[{cname}] Unexpected error: {e}") + + +def update_entrypoint_install_path(script_path: str | Path, new_path: str) -> str: + """Update the micromamba pip install path in the given bash script.""" + script_path = Path(script_path) + lines = script_path.read_text().splitlines() + + new_lines = [] + for line in lines: + if line.strip().startswith("micromamba run -n default_user_env pip install -e"): + line = f'micromamba run -n default_user_env pip install -e {new_path}' + new_lines.append(line) + + updated_text = "\n".join(new_lines) + "\n" + + script_path.write_text(updated_text) + + return updated_text From 0a5ab966881172732b2eaa8988e67a22ae89c762 Mon Sep 17 00:00:00 2001 From: Yogesh Kumar Baljeet Singh Date: Tue, 16 Sep 2025 09:34:54 +0200 Subject: [PATCH 07/15] ruff check --- src/gaiaflow/cli/commands/minikube.py | 8 +++----- src/gaiaflow/cli/commands/mlops.py | 3 +-- src/gaiaflow/managers/mlops_manager.py | 12 +++++++----- src/gaiaflow/managers/utils.py | 2 +- 4 files changed, 12 insertions(+), 13 deletions(-) diff --git a/src/gaiaflow/cli/commands/minikube.py b/src/gaiaflow/cli/commands/minikube.py index 1979bb2..c609e74 100644 --- a/src/gaiaflow/cli/commands/minikube.py +++ b/src/gaiaflow/cli/commands/minikube.py @@ -58,8 +58,7 @@ def start( @app.command(help="Stop Gaiaflow production-like services.") -def stop( -): +def stop(): imports = load_imports() project_path = Path.cwd() gaiaflow_path, user_project_path = imports.create_gaiaflow_context_path( @@ -99,7 +98,7 @@ def restart( gaiaflow_path=gaiaflow_path, user_project_path=user_project_path, action=imports.BaseAction.RESTART, - force_new=force_new + force_new=force_new, ) @@ -133,8 +132,7 @@ def dockerize( help="Create a config file for Airflow to talk to Kubernetes " "cluster. To be used only when debugging required." ) -def create_config( -): +def create_config(): imports = load_imports() project_path = Path.cwd() gaiaflow_path, user_project_path = imports.create_gaiaflow_context_path( diff --git a/src/gaiaflow/cli/commands/mlops.py b/src/gaiaflow/cli/commands/mlops.py index 3b68b78..9bf0dbc 100644 --- a/src/gaiaflow/cli/commands/mlops.py +++ b/src/gaiaflow/cli/commands/mlops.py @@ -289,8 +289,7 @@ def dockerize( "this, as the container environments are updated based on " "its contents." ) -def update_deps( -): +def update_deps(): imports = load_imports() project_path = Path.cwd() gaiaflow_path, user_project_path = imports.create_gaiaflow_context_path( diff --git a/src/gaiaflow/managers/mlops_manager.py b/src/gaiaflow/managers/mlops_manager.py index 69198f7..68de912 100644 --- a/src/gaiaflow/managers/mlops_manager.py +++ b/src/gaiaflow/managers/mlops_manager.py @@ -28,7 +28,6 @@ create_directory, delete_project_state, env_exists, - find_python_packages, gaiaflow_path_exists_in_state, handle_error, log_error, @@ -36,8 +35,8 @@ run, save_project_state, set_permissions, - update_micromamba_env_in_docker, update_entrypoint_install_path, + update_micromamba_env_in_docker, ) @@ -427,7 +426,9 @@ def _collect_volumes(self, compose_data: dict) -> list[str]: # their package set_permissions(self.user_project_path, 0o755) - new_volumes.append(f"{self.user_project_path.resolve().as_posix()}:/opt/airflow/{self.user_project_path.name}") + new_volumes.append( + f"{self.user_project_path.resolve().as_posix()}:/opt/airflow/{self.user_project_path.name}" + ) # Add special mounts for prod_local mode kube_config = ( @@ -472,8 +473,9 @@ def _update_files(self): entrypoint_path = ( self.gaiaflow_path / "_docker" / "docker-compose" / "entrypoint.sh" ) - update_entrypoint_install_path(entrypoint_path, - str(self.user_project_path.name)) + update_entrypoint_install_path( + entrypoint_path, str(self.user_project_path.name) + ) set_permissions(entrypoint_path) convert_crlf_to_lf(entrypoint_path) diff --git a/src/gaiaflow/managers/utils.py b/src/gaiaflow/managers/utils.py index 3fcd0a7..7905067 100644 --- a/src/gaiaflow/managers/utils.py +++ b/src/gaiaflow/managers/utils.py @@ -316,7 +316,7 @@ def update_entrypoint_install_path(script_path: str | Path, new_path: str) -> st new_lines = [] for line in lines: if line.strip().startswith("micromamba run -n default_user_env pip install -e"): - line = f'micromamba run -n default_user_env pip install -e {new_path}' + line = f"micromamba run -n default_user_env pip install -e {new_path}" new_lines.append(line) updated_text = "\n".join(new_lines) + "\n" From 05c42abf0a0b337cd5f47a34e1a27a824940941f Mon Sep 17 00:00:00 2001 From: Yogesh Kumar Baljeet Singh Date: Tue, 16 Sep 2025 12:39:48 +0200 Subject: [PATCH 08/15] add docker-proxy to avoid var/run/docker.sock permission issues and add op_kwargs to create_task --- src/_docker/docker-compose/docker-compose.yml | 11 ++++++++++ src/gaiaflow/constants.py | 1 + src/gaiaflow/core/create_task.py | 2 ++ src/gaiaflow/core/operators.py | 20 ++++++++++++++++--- 4 files changed, 31 insertions(+), 3 deletions(-) diff --git a/src/_docker/docker-compose/docker-compose.yml b/src/_docker/docker-compose/docker-compose.yml index d4b203d..087ba51 100644 --- a/src/_docker/docker-compose/docker-compose.yml +++ b/src/_docker/docker-compose/docker-compose.yml @@ -54,6 +54,17 @@ x-airflow-common: - ml-network services: + docker-proxy: + image: alpine/socat + container_name: docker-proxy + command: "TCP4-LISTEN:2375,fork,reuseaddr UNIX-CONNECT:/var/run/docker.sock" + ports: + - "2376:2375" + volumes: + - /var/run/docker.sock:/var/run/docker.sock + networks: + - ml-network + postgres-mlflow: image: postgres:13 diff --git a/src/gaiaflow/constants.py b/src/gaiaflow/constants.py index 82c843a..e5beeee 100644 --- a/src/gaiaflow/constants.py +++ b/src/gaiaflow/constants.py @@ -42,6 +42,7 @@ class Service(str, Enum): "airflow-dag-processor", "airflow-triggerer", "postgres-airflow", + "docker-proxy", ] MLFLOW_SERVICES = ["mlflow", "postgres-mlflow"] diff --git a/src/gaiaflow/core/create_task.py b/src/gaiaflow/core/create_task.py index 7dd4f7a..c28dd8a 100755 --- a/src/gaiaflow/core/create_task.py +++ b/src/gaiaflow/core/create_task.py @@ -34,6 +34,7 @@ def create_task( env_vars: dict | None = None, retries: int = 3, dag=None, + **op_kwargs ): """It is a high-level abstraction on top of Apache Airflow operators. @@ -69,6 +70,7 @@ def create_task( retries=retries, params=dag_params, mode=gaiaflow_mode, + **op_kwargs ) return operator.create_task() diff --git a/src/gaiaflow/core/operators.py b/src/gaiaflow/core/operators.py index f481865..abea1cb 100644 --- a/src/gaiaflow/core/operators.py +++ b/src/gaiaflow/core/operators.py @@ -60,6 +60,7 @@ def __init__( retries: int, params: dict, mode: str, + **op_kwargs, ): self.task_id = task_id self.func_path = func_path @@ -69,6 +70,7 @@ def __init__( self.retries = retries self.params = params self.mode = mode + self.op_kwargs = op_kwargs ( self.func_args, @@ -156,6 +158,7 @@ def run_wrapper(**op_kwargs): retries=self.retries, expect_airflow=False, expect_pendulum=False, + **self.op_kwargs ) @@ -222,10 +225,14 @@ def create_task(self): # }, # ) + command = self.op_kwargs.pop("cmds", None) + if command is None: + command = ["python", "-m", "runner"] + return KubernetesPodOperator( task_id=self.task_id, image=self.image, - cmds=["python", "-m", "runner"], + cmds=command, env_vars=all_env_vars, env_from=env_from, get_logs=True, @@ -235,6 +242,7 @@ def create_task(self): do_xcom_push=True, retries=self.retries, params=self.params, + **self.op_kwargs, # container_resources=resources, ) @@ -282,6 +290,10 @@ def create_task(self): safe_image_name = self.image.replace(":", "_").replace("/", "_") + command = self.op_kwargs.pop("command", None) + if command is None: + command = ["python", "-m", "runner"] + return DockerOperator( task_id=self.task_id, image=self.image, @@ -293,8 +305,9 @@ def create_task(self): + "_container", api_version="auto", auto_remove="success", - command=["python", "-m", "runner"], - docker_url="unix://var/run/docker.sock", + command=command, + # docker_url="unix://var/run/docker.sock", + docker_url='tcp://docker-proxy:2375', # docker_url="tcp://host.docker.internal:2375", environment=combined_env, network_mode="docker-compose_ml-network", @@ -303,4 +316,5 @@ def create_task(self): retrieve_output=True, retrieve_output_path="/tmp/script.out", xcom_all=False, + **self.op_kwargs, ) From 437a1f814c293e3a7f3628e37c82f85db2909680 Mon Sep 17 00:00:00 2001 From: Yogesh Kumar Baljeet Singh Date: Tue, 16 Sep 2025 12:40:27 +0200 Subject: [PATCH 09/15] update image and container names --- src/gaiaflow/managers/mlops_manager.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/gaiaflow/managers/mlops_manager.py b/src/gaiaflow/managers/mlops_manager.py index 68de912..3b8eb16 100644 --- a/src/gaiaflow/managers/mlops_manager.py +++ b/src/gaiaflow/managers/mlops_manager.py @@ -51,6 +51,7 @@ class DockerResources: "minio/mc:latest", "minio/minio:latest", "postgres:13", + "alpine/socat", ] AIRFLOW_CONTAINERS = [ @@ -58,6 +59,7 @@ class DockerResources: "airflow-scheduler", "airflow-dag-processor", "airflow-triggerer", + "docker-proxy" ] VOLUMES = [ From 055c6a8da65f4447970143d0a3f205f055096429 Mon Sep 17 00:00:00 2001 From: Yogesh Kumar Baljeet Singh Date: Tue, 16 Sep 2025 12:41:15 +0200 Subject: [PATCH 10/15] implement Strategy pattern for DockerBuilder and also support user-provide dockerfile for build --- src/gaiaflow/cli/commands/minikube.py | 11 +- src/gaiaflow/cli/commands/mlops.py | 11 +- src/gaiaflow/managers/minikube_manager.py | 180 +++++++++++++--------- 3 files changed, 129 insertions(+), 73 deletions(-) diff --git a/src/gaiaflow/cli/commands/minikube.py b/src/gaiaflow/cli/commands/minikube.py index c609e74..4bdd659 100644 --- a/src/gaiaflow/cli/commands/minikube.py +++ b/src/gaiaflow/cli/commands/minikube.py @@ -109,6 +109,10 @@ def dockerize( image_name: str = typer.Option( DEFAULT_IMAGE_NAME, "--image-name", "-i", help=("Name of your image.") ), + dockerfile_path: Path = typer.Option( + None, "--dockerfile-path", "-d", help=("Path to your custom " + "Dockerfile") + ), ): imports = load_imports() project_path = Path.cwd() @@ -119,12 +123,17 @@ def dockerize( if not gaiaflow_path_exists: typer.echo("Please create a project with Gaiaflow before running this command.") return + if dockerfile_path: + docker_build_mode = "minikube-user" + else: + docker_build_mode = "minikube" imports.MinikubeManager.run( gaiaflow_path=gaiaflow_path, user_project_path=user_project_path, action=imports.ExtendedAction.DOCKERIZE, - local=False, + docker_build_mode=docker_build_mode, image_name=image_name, + dockerfile_path=dockerfile_path, ) diff --git a/src/gaiaflow/cli/commands/mlops.py b/src/gaiaflow/cli/commands/mlops.py index 9bf0dbc..efbebc8 100644 --- a/src/gaiaflow/cli/commands/mlops.py +++ b/src/gaiaflow/cli/commands/mlops.py @@ -255,6 +255,10 @@ def dockerize( image_name: str = typer.Option( DEFAULT_IMAGE_NAME, "--image-name", "-i", help=("Name of your image.") ), + dockerfile_path: Path = typer.Option( + None, "--dockerfile-path", "-d", help=("Path to your custom " + "Dockerfile") + ), ): imports = load_imports() project_path = Path.cwd() @@ -270,13 +274,16 @@ def dockerize( f"Gaiaflow project already exists at {gaiaflow_path}. Skipping " f"saving to the state" ) - + if dockerfile_path: + docker_build_mode = "local-user" + else: + docker_build_mode = "local" typer.echo("Running dockerize") imports.MinikubeManager.run( gaiaflow_path=gaiaflow_path, user_project_path=user_project_path, action=imports.ExtendedAction.DOCKERIZE, - local=True, + docker_build_mode=docker_build_mode, image_name=image_name, ) diff --git a/src/gaiaflow/managers/minikube_manager.py b/src/gaiaflow/managers/minikube_manager.py index 4964627..03c43ce 100644 --- a/src/gaiaflow/managers/minikube_manager.py +++ b/src/gaiaflow/managers/minikube_manager.py @@ -4,7 +4,7 @@ import subprocess from contextlib import contextmanager from pathlib import Path -from typing import Any, Set +from typing import Any, Set, Literal import yaml @@ -95,77 +95,110 @@ def run_cmd(self, args: list[str], **kwargs): return subprocess.run(full_cmd, **kwargs) +class BaseDockerBuilder: + """Abstract docker builder with optional hooks.""" + + @classmethod + def get_docker_builder(cls, mode: str, **kwargs): + builder_cls = BUILDER_REGISTRY.get(mode) + if not builder_cls: + raise ValueError(f"Unknown Docker build mode: {mode}") + return builder_cls(**kwargs) + + def pre_build(self, image_name: str, dockerfile_path: Path, project_path: Path): + """Override this if you want a different or no pre_build""" + log_info(f"Updating Dockerfile at {dockerfile_path}") + DockerHelper._add_copy_statements_to_dockerfile( + str(dockerfile_path), find_python_packages(project_path) + ) + runner_src = Path(__file__).parent.parent.resolve() / "core" / "runner.py" + runner_dest = project_path / "runner.py" + return temporary_copy(runner_src, runner_dest) + + def build(self, image_name: str, dockerfile_path: Path, project_path: Path): + raise NotImplementedError + + def post_build(self, image_name: str, dockerfile_path: Path, project_path: Path): + pass + +class LocalDockerBuilder(BaseDockerBuilder): + def build(self, image_name: str, dockerfile_path: Path, project_path: Path): + log_info(f"Building Docker image [{image_name}] locally") + run( + ["docker", "build", "-t", image_name, "-f", str(dockerfile_path), str(project_path)], + "Error building Docker image locally", + ) + +class MinikubeDockerBuilder(BaseDockerBuilder): + def __init__(self, minikube_helper: MinikubeHelper): + self.minikube_helper = minikube_helper + + def build(self, image_name: str, dockerfile_path: Path, project_path: Path): + log_info(f"Building Docker image [{image_name}] in Minikube context") + result = self.minikube_helper.run_cmd( + ["docker-env", "--shell", "bash"], stdout=subprocess.PIPE, check=True + ) + env = DockerHelper._parse_minikube_env(result.stdout.decode()) + run( + ["docker", "build", "-t", image_name, "-f", str(dockerfile_path), str(project_path)], + "Error building Docker image inside Minikube", + env=env, + ) + +class LocalUserCustomImageDockerBuilder(LocalDockerBuilder): + def pre_build(self, image_name: str, dockerfile_path: Path, project_path: Path): + pass + def build(self, image_name: str, dockerfile_path: Path, project_path: Path): + log_info("Building user provided dockerfile") + super().build(image_name, dockerfile_path, project_path) + +class MinikubeUserCustomImageDockerBuilder(MinikubeDockerBuilder): + def pre_build(self, image_name: str, dockerfile_path: Path, project_path: Path): + pass + + def build(self, image_name: str, dockerfile_path: Path, project_path: Path): + log_info("Building user provided dockerfile") + super().build(image_name, dockerfile_path, project_path) + +BUILDER_REGISTRY = { + "local": LocalDockerBuilder, + "minikube": MinikubeDockerBuilder, + "local-user": LocalUserCustomImageDockerBuilder, + "minikube-user": MinikubeUserCustomImageDockerBuilder +} + class DockerHelper: def __init__( self, image_name: str, project_path: Path, - local: bool, - minikube_helper: MinikubeHelper, + builder: BaseDockerBuilder ): self.image_name = image_name self.project_path = project_path - self.local = local - self.minikube_helper = minikube_helper + self.builder = builder 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" + pre_build_ctx = self.builder.pre_build( + self.image_name, dockerfile_path, self.project_path + ) + if pre_build_ctx: + with pre_build_ctx: + self.builder.build(self.image_name, dockerfile_path, self.project_path) + else: + self.builder.build(self.image_name, dockerfile_path, self.project_path) - with temporary_copy(runner_src, runner_dest): - if self.local: - self._build_local(dockerfile_path) - else: - self._build_minikube(dockerfile_path) + self.builder.post_build(self.image_name, dockerfile_path, self.project_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() @@ -225,17 +258,18 @@ def __init__(self, gaiaflow_path: Path, os_type: str): self.os_type = os_type def create_inline(self): - kube_config = Path.home() / ".kube" / "config" - backup_config = kube_config.with_suffix(".backup") + if self.os_type == "linux" or is_wsl(): + 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) + 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.") + if 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(): @@ -289,6 +323,7 @@ def _write_inline(self, kube_config: Path): class MinikubeManager(BaseGaiaflowManager): + allowed_kwargs = {"secret_name", "secret_data", "dockerfile_path"} def __init__( self, gaiaflow_path: Path, @@ -296,26 +331,29 @@ def __init__( action: Action, force_new: bool = False, prune: bool = False, - local: bool = False, + docker_build_mode: Literal["local", "minikube"] = "local", image_name: str = "", **kwargs, ): - # if kwargs: - # raise TypeError(f"Unexpected keyword arguments: {list(kwargs.keys())}") + if kwargs: + for key in kwargs: + if key not in self.allowed_kwargs: + raise TypeError(f"Unexpected keyword argument: {key}") + self.minikube_profile = "airflow" # TODO: get the docker image name automatically # For CI, get the package name, version and create repository. See # in test-airflow-ci test_ecr_push.yml self.os_type = platform.system().lower() - self.local = local self.image_name = image_name self.minikube_helper = MinikubeHelper() + builder = BaseDockerBuilder.get_docker_builder(docker_build_mode, + minikube_helper=self.minikube_helper) self.docker_helper = DockerHelper( image_name=image_name, project_path=user_project_path, - local=local, - minikube_helper=self.minikube_helper, + builder=builder, ) self.kube_helper = KubeConfigHelper( gaiaflow_path=gaiaflow_path, os_type=self.os_type @@ -338,7 +376,7 @@ def _get_valid_actions(self) -> Set[Action]: @classmethod def run(cls, **kwargs): - action = kwargs.get("action") + action = kwargs.get("action", None) if action is None: raise ValueError("Missing required argument 'action'") @@ -349,7 +387,8 @@ def run(cls, **kwargs): BaseAction.STOP: manager.stop, BaseAction.RESTART: manager.restart, BaseAction.CLEANUP: manager.cleanup, - ExtendedAction.DOCKERIZE: manager.build_docker_image, + ExtendedAction.DOCKERIZE: lambda: manager.build_docker_image( + kwargs["dockerfile_path"]), ExtendedAction.CREATE_CONFIG: manager.create_kube_config_inline, ExtendedAction.CREATE_SECRET: lambda: manager.create_secrets( kwargs["secret_name"], kwargs["secret_data"] @@ -391,8 +430,9 @@ def stop(self): def create_kube_config_inline(self): self.kube_helper.create_inline() - def build_docker_image(self): - dockerfile_path = self.gaiaflow_path / "_docker" / "user-package" / "Dockerfile" + def build_docker_image(self, dockerfile_path: str): + if not dockerfile_path: + dockerfile_path = self.gaiaflow_path / "_docker" / "user-package" / "Dockerfile" self.docker_helper.build_image(dockerfile_path) def create_secrets(self, secret_name: str, secret_data: dict[str, Any]): From ffc868998bac45561fdd8ae6255384d2d556f8d5 Mon Sep 17 00:00:00 2001 From: Yogesh Kumar Baljeet Singh Date: Tue, 16 Sep 2025 13:10:34 +0200 Subject: [PATCH 11/15] rename builder to handler and add few more docker operations --- src/gaiaflow/managers/minikube_manager.py | 198 +++++++++++++--------- src/gaiaflow/managers/utils.py | 4 +- 2 files changed, 123 insertions(+), 79 deletions(-) diff --git a/src/gaiaflow/managers/minikube_manager.py b/src/gaiaflow/managers/minikube_manager.py index 03c43ce..7b573f7 100644 --- a/src/gaiaflow/managers/minikube_manager.py +++ b/src/gaiaflow/managers/minikube_manager.py @@ -95,20 +95,20 @@ def run_cmd(self, args: list[str], **kwargs): return subprocess.run(full_cmd, **kwargs) -class BaseDockerBuilder: - """Abstract docker builder with optional hooks.""" +class BaseDockerHandler: + """Abstract docker handler with optional hooks.""" @classmethod - def get_docker_builder(cls, mode: str, **kwargs): - builder_cls = BUILDER_REGISTRY.get(mode) - if not builder_cls: + def get_docker_handler(cls, mode: str, **kwargs): + handler_cls = HANDLER_REGISTRY.get(mode) + if not handler_cls: raise ValueError(f"Unknown Docker build mode: {mode}") - return builder_cls(**kwargs) + return handler_cls(**kwargs) def pre_build(self, image_name: str, dockerfile_path: Path, project_path: Path): """Override this if you want a different or no pre_build""" log_info(f"Updating Dockerfile at {dockerfile_path}") - DockerHelper._add_copy_statements_to_dockerfile( + BaseDockerHandler._add_copy_statements_to_dockerfile( str(dockerfile_path), find_python_packages(project_path) ) runner_src = Path(__file__).parent.parent.resolve() / "core" / "runner.py" @@ -121,7 +121,61 @@ def build(self, image_name: str, dockerfile_path: Path, project_path: Path): def post_build(self, image_name: str, dockerfile_path: Path, project_path: Path): pass -class LocalDockerBuilder(BaseDockerBuilder): + def list_images(self): + raise NotImplementedError + + def remove_image(self, image_name: str): + raise NotImplementedError + + def prune_images(self): + raise NotImplementedError + + def _update_dockerfile(self, dockerfile_path: Path): + BaseDockerHandler._add_copy_statements_to_dockerfile( + str(dockerfile_path), find_python_packages(self.project_path) + ) + + @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 LocalDockerHandler(BaseDockerHandler): def build(self, image_name: str, dockerfile_path: Path, project_path: Path): log_info(f"Building Docker image [{image_name}] locally") run( @@ -129,30 +183,68 @@ def build(self, image_name: str, dockerfile_path: Path, project_path: Path): "Error building Docker image locally", ) -class MinikubeDockerBuilder(BaseDockerBuilder): + def list_images(self): + run(["docker", "image", "ls"], "Error listing Docker images locally") + + def remove_image(self, image_name: str): + run(["docker", "rmi", "-f", image_name], f"Error removing Docker image {image_name} " + "locally") + + def prune_images(self): + run(["docker", "image", "prune", "-f"], "Error pruning Docker images " + "locally") + +class MinikubeDockerHandler(BaseDockerHandler): def __init__(self, minikube_helper: MinikubeHelper): self.minikube_helper = minikube_helper + self.env = self._get_minikube_env() def build(self, image_name: str, dockerfile_path: Path, project_path: Path): log_info(f"Building Docker image [{image_name}] in Minikube context") - result = self.minikube_helper.run_cmd( - ["docker-env", "--shell", "bash"], stdout=subprocess.PIPE, check=True - ) - env = DockerHelper._parse_minikube_env(result.stdout.decode()) run( ["docker", "build", "-t", image_name, "-f", str(dockerfile_path), str(project_path)], "Error building Docker image inside Minikube", - env=env, + env=self.env, + ) + + def list_images(self): + run(["docker", "image", "ls"],"Error listing Docker images inside Minikube", env=self.env) + + def remove_image(self, image_name: str): + run(["docker", "rmi", "-f", image_name], f"Error removing Docker image {image_name} " + "inside Minikube", env=self.env) + + def prune_images(self): + run(["docker", "image", "prune", "-f"], "Error pruning Docker images " + "inside Minikube", env=self.env) + + def _get_minikube_env(self): + result = self.minikube_helper.run_cmd( + ["docker-env", "--shell", "bash"], stdout=subprocess.PIPE, check=True ) + return MinikubeDockerHandler._parse_minikube_env(result.stdout.decode()) -class LocalUserCustomImageDockerBuilder(LocalDockerBuilder): + @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 + +class LocalUserCustomImageDockerHandler(LocalDockerHandler): def pre_build(self, image_name: str, dockerfile_path: Path, project_path: Path): pass + def build(self, image_name: str, dockerfile_path: Path, project_path: Path): log_info("Building user provided dockerfile") super().build(image_name, dockerfile_path, project_path) -class MinikubeUserCustomImageDockerBuilder(MinikubeDockerBuilder): +class MinikubeUserCustomImageDockerHandler(MinikubeDockerHandler): def pre_build(self, image_name: str, dockerfile_path: Path, project_path: Path): pass @@ -160,11 +252,11 @@ def build(self, image_name: str, dockerfile_path: Path, project_path: Path): log_info("Building user provided dockerfile") super().build(image_name, dockerfile_path, project_path) -BUILDER_REGISTRY = { - "local": LocalDockerBuilder, - "minikube": MinikubeDockerBuilder, - "local-user": LocalUserCustomImageDockerBuilder, - "minikube-user": MinikubeUserCustomImageDockerBuilder +HANDLER_REGISTRY = { + "local": LocalDockerHandler, + "minikube": MinikubeDockerHandler, + "local-user": LocalUserCustomImageDockerHandler, + "minikube-user": MinikubeUserCustomImageDockerHandler } class DockerHelper: @@ -172,7 +264,7 @@ def __init__( self, image_name: str, project_path: Path, - builder: BaseDockerBuilder + builder: BaseDockerHandler ): self.image_name = image_name self.project_path = project_path @@ -194,62 +286,14 @@ def build_image(self, dockerfile_path: Path): self.builder.post_build(self.image_name, dockerfile_path, self.project_path) - def _update_dockerfile(self, dockerfile_path: Path): - DockerHelper._add_copy_statements_to_dockerfile( - str(dockerfile_path), find_python_packages(self.project_path) - ) - - @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") + def list_images(self): + self.builder.list_images() - updated_lines = ( - lines[: env_index + 1] - + copy_lines # - + lines[entrypoint_index:] - ) - with open(dockerfile_path, "w") as f: - f.writelines(updated_lines) + def remove_image(self, image_name: str): + self.builder.remove_image(image_name) - print("Dockerfile updated with COPY statements.") + def prune_images(self): + self.builder.prune_images() class KubeConfigHelper: @@ -348,7 +392,7 @@ def __init__( self.image_name = image_name self.minikube_helper = MinikubeHelper() - builder = BaseDockerBuilder.get_docker_builder(docker_build_mode, + builder = BaseDockerHandler.get_docker_builder(docker_build_mode, minikube_helper=self.minikube_helper) self.docker_helper = DockerHelper( image_name=image_name, diff --git a/src/gaiaflow/managers/utils.py b/src/gaiaflow/managers/utils.py index 7905067..ec7cdd6 100644 --- a/src/gaiaflow/managers/utils.py +++ b/src/gaiaflow/managers/utils.py @@ -50,9 +50,9 @@ def log_error(message: str): ) -def run(command: list, error_message: str, env=None): +def run(command: list, error_message: str, **kwargs): try: - subprocess.call(command, env=env) + subprocess.call(command, **kwargs) except Exception: log_error(error_message) raise From 0de184165ec394424b18b2beccfb1c7b57c31c29 Mon Sep 17 00:00:00 2001 From: Yogesh Kumar Baljeet Singh Date: Tue, 16 Sep 2025 14:31:28 +0200 Subject: [PATCH 12/15] moved helper classes to helper --- src/gaiaflow/cli/commands/minikube.py | 45 +- src/gaiaflow/cli/commands/mlops.py | 45 +- src/gaiaflow/constants.py | 2 + src/gaiaflow/core/create_task.py | 4 +- src/gaiaflow/core/operators.py | 4 +- src/gaiaflow/managers/helpers.py | 540 ++++++++++++++++++++++ src/gaiaflow/managers/minikube_manager.py | 392 ++-------------- src/gaiaflow/managers/mlops_manager.py | 164 +------ src/gaiaflow/managers/utils.py | 1 - 9 files changed, 671 insertions(+), 526 deletions(-) create mode 100644 src/gaiaflow/managers/helpers.py diff --git a/src/gaiaflow/cli/commands/minikube.py b/src/gaiaflow/cli/commands/minikube.py index 4bdd659..33449d8 100644 --- a/src/gaiaflow/cli/commands/minikube.py +++ b/src/gaiaflow/cli/commands/minikube.py @@ -5,6 +5,7 @@ import typer from gaiaflow.constants import DEFAULT_IMAGE_NAME +from gaiaflow.managers.helpers import DockerHandlerMode app = typer.Typer() fs = fsspec.filesystem("file") @@ -110,8 +111,7 @@ def dockerize( DEFAULT_IMAGE_NAME, "--image-name", "-i", help=("Name of your image.") ), dockerfile_path: Path = typer.Option( - None, "--dockerfile-path", "-d", help=("Path to your custom " - "Dockerfile") + None, "--dockerfile-path", "-d", help=("Path to your custom Dockerfile") ), ): imports = load_imports() @@ -137,6 +137,47 @@ def dockerize( ) +@app.command(help="List all the docker images in your system") +def list_images(): + imports = load_imports() + project_path = Path.cwd() + gaiaflow_path, user_project_path = imports.create_gaiaflow_context_path( + project_path + ) + gaiaflow_path_exists = imports.gaiaflow_path_exists_in_state(gaiaflow_path, True) + if not gaiaflow_path_exists: + typer.echo("Please create a project with Gaiaflow before running this command.") + return + imports.MinikubeManager.run( + gaiaflow_path=gaiaflow_path, + user_project_path=user_project_path, + action=imports.ExtendedAction.LIST_IMAGES, + docker_handler_mode=DockerHandlerMode.MINIKUBE, + ) + +@app.command(help="Delete a docker image from your system") +def remove_image(image_name: str = typer.Option( + DEFAULT_IMAGE_NAME, "--image-name", "-i", help=("Name of image " + "to be deleted.") + ),): + imports = load_imports() + project_path = Path.cwd() + gaiaflow_path, user_project_path = imports.create_gaiaflow_context_path( + project_path + ) + gaiaflow_path_exists = imports.gaiaflow_path_exists_in_state(gaiaflow_path, True) + if not gaiaflow_path_exists: + typer.echo("Please create a project with Gaiaflow before running this command.") + return + imports.MinikubeManager.run( + gaiaflow_path=gaiaflow_path, + user_project_path=user_project_path, + action=imports.ExtendedAction.REMOVE_IMAGE, + image_name=image_name, + docker_handler_mode=DockerHandlerMode.MINIKUBE, + ) + + @app.command( help="Create a config file for Airflow to talk to Kubernetes " "cluster. To be used only when debugging required." diff --git a/src/gaiaflow/cli/commands/mlops.py b/src/gaiaflow/cli/commands/mlops.py index efbebc8..5415543 100644 --- a/src/gaiaflow/cli/commands/mlops.py +++ b/src/gaiaflow/cli/commands/mlops.py @@ -6,6 +6,7 @@ import typer from gaiaflow.constants import DEFAULT_IMAGE_NAME, Service +from gaiaflow.managers.helpers import DockerHandlerMode app = typer.Typer() fs = fsspec.filesystem("file") @@ -256,8 +257,7 @@ def dockerize( DEFAULT_IMAGE_NAME, "--image-name", "-i", help=("Name of your image.") ), dockerfile_path: Path = typer.Option( - None, "--dockerfile-path", "-d", help=("Path to your custom " - "Dockerfile") + None, "--dockerfile-path", "-d", help=("Path to your custom Dockerfile") ), ): imports = load_imports() @@ -288,6 +288,47 @@ def dockerize( ) + +@app.command(help="List all the docker images in your system") +def list_images(): + imports = load_imports() + project_path = Path.cwd() + gaiaflow_path, user_project_path = imports.create_gaiaflow_context_path( + project_path + ) + gaiaflow_path_exists = imports.gaiaflow_path_exists_in_state(gaiaflow_path, True) + if not gaiaflow_path_exists: + typer.echo("Please create a project with Gaiaflow before running this command.") + return + imports.MinikubeManager.run( + gaiaflow_path=gaiaflow_path, + user_project_path=user_project_path, + action=imports.ExtendedAction.LIST_IMAGES, + docker_handler_mode=DockerHandlerMode.LOCAL, + ) + +@app.command(help="Delete a docker image from your system") +def remove_image(image_name: str = typer.Option( + DEFAULT_IMAGE_NAME, "--image-name", "-i", help=("Name of image " + "to be deleted.") + ),): + imports = load_imports() + project_path = Path.cwd() + gaiaflow_path, user_project_path = imports.create_gaiaflow_context_path( + project_path + ) + gaiaflow_path_exists = imports.gaiaflow_path_exists_in_state(gaiaflow_path, True) + if not gaiaflow_path_exists: + typer.echo("Please create a project with Gaiaflow before running this command.") + return + imports.MinikubeManager.run( + gaiaflow_path=gaiaflow_path, + user_project_path=user_project_path, + action=imports.ExtendedAction.REMOVE_IMAGE, + image_name=image_name, + docker_handler_mode=DockerHandlerMode.LOCAL, + ) + @app.command( help="Update the dependencies for the Airflow tasks. This command " "synchronizes the running container environments with the project's" diff --git a/src/gaiaflow/constants.py b/src/gaiaflow/constants.py index e5beeee..8898f18 100644 --- a/src/gaiaflow/constants.py +++ b/src/gaiaflow/constants.py @@ -20,6 +20,8 @@ class ExtendedAction: CREATE_CONFIG = Action("create_config") CREATE_SECRET = Action("create_secret") UPDATE_DEPS = Action("update_deps") + LIST_IMAGES = Action("list-images") + REMOVE_IMAGE = Action("remove-image") GAIAFLOW_CONFIG_DIR = Path.home() / ".gaiaflow" diff --git a/src/gaiaflow/core/create_task.py b/src/gaiaflow/core/create_task.py index c28dd8a..d3e598b 100755 --- a/src/gaiaflow/core/create_task.py +++ b/src/gaiaflow/core/create_task.py @@ -34,7 +34,7 @@ def create_task( env_vars: dict | None = None, retries: int = 3, dag=None, - **op_kwargs + **op_kwargs, ): """It is a high-level abstraction on top of Apache Airflow operators. @@ -70,7 +70,7 @@ def create_task( retries=retries, params=dag_params, mode=gaiaflow_mode, - **op_kwargs + **op_kwargs, ) return operator.create_task() diff --git a/src/gaiaflow/core/operators.py b/src/gaiaflow/core/operators.py index abea1cb..54b575f 100644 --- a/src/gaiaflow/core/operators.py +++ b/src/gaiaflow/core/operators.py @@ -158,7 +158,7 @@ def run_wrapper(**op_kwargs): retries=self.retries, expect_airflow=False, expect_pendulum=False, - **self.op_kwargs + **self.op_kwargs, ) @@ -307,7 +307,7 @@ def create_task(self): auto_remove="success", command=command, # docker_url="unix://var/run/docker.sock", - docker_url='tcp://docker-proxy:2375', + docker_url="tcp://docker-proxy:2375", # docker_url="tcp://host.docker.internal:2375", environment=combined_env, network_mode="docker-compose_ml-network", diff --git a/src/gaiaflow/managers/helpers.py b/src/gaiaflow/managers/helpers.py new file mode 100644 index 0000000..3fffb57 --- /dev/null +++ b/src/gaiaflow/managers/helpers.py @@ -0,0 +1,540 @@ +import os +import shutil +import socket +import subprocess +from contextlib import contextmanager +from pathlib import Path + +import psutil +import yaml + +from gaiaflow.constants import AIRFLOW_SERVICES, MINIO_SERVICES, MLFLOW_SERVICES +from gaiaflow.managers.utils import ( + env_exists, + find_python_packages, + handle_error, + is_wsl, + log_error, + log_info, + run, +) + + +@contextmanager +def temporary_copy(src: Path, dest: Path): + print("copying...", src, dest) + shutil.copyfile(src, dest) + try: + yield + finally: + if dest.exists(): + dest.unlink() + + +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", + "alpine/socat", + ] + + AIRFLOW_CONTAINERS = [ + "airflow-apiserver", + "airflow-scheduler", + "airflow-dag-processor", + "airflow-triggerer", + "docker-proxy", + ] + + VOLUMES = [ + "docker-compose_postgres-db-volume-airflow", + "docker-compose_postgres-db-volume-mlflow", + ] + + SERVICES = { + "airflow": AIRFLOW_SERVICES, + "mlflow": MLFLOW_SERVICES, + "minio": MINIO_SERVICES, + } + + +class DockerComposeHelper: + 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/docker-compose/docker-compose.yml", + ] + if self.is_prod_local: + base += [ + "-f", + f"{self.gaiaflow_path}/_docker/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 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 DockerHandlerMode: + LOCAL = "local" + MINIKUBE = "minikube" + LOCAL_USER = "local-user" + MINIKUBE_USER = "minikube-user" + +class BaseDockerHandler: + """Abstract docker handler with optional hooks.""" + def __init__(self, **kwargs): + pass + + @classmethod + def get_docker_handler(cls, mode: DockerHandlerMode, **kwargs): + handler_cls = HANDLER_REGISTRY.get(mode) + if not handler_cls: + raise ValueError(f"Unknown Docker build mode: {mode}") + return handler_cls(**kwargs) + + def pre_build(self, image_name: str, dockerfile_path: Path, project_path: Path): + """Override this if you want a different or no pre_build""" + log_info(f"Updating Dockerfile at {dockerfile_path}") + BaseDockerHandler._add_copy_statements_to_dockerfile( + str(dockerfile_path), find_python_packages(project_path) + ) + runner_src = Path(__file__).parent.parent.resolve() / "core" / "runner.py" + runner_dest = project_path / "runner.py" + return temporary_copy(runner_src, runner_dest) + + def build(self, image_name: str, dockerfile_path: Path, project_path: Path): + raise NotImplementedError + + def post_build(self, image_name: str, dockerfile_path: Path, project_path: Path): + pass + + def list_images(self): + raise NotImplementedError + + def remove_image(self, image_name: str): + raise NotImplementedError + + def _update_dockerfile(self, dockerfile_path: Path): + BaseDockerHandler._add_copy_statements_to_dockerfile( + str(dockerfile_path), find_python_packages(self.project_path) + ) + + @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 LocalDockerHandler(BaseDockerHandler): + def build(self, image_name: str, dockerfile_path: Path, project_path: Path): + log_info(f"Building Docker image [{image_name}] locally") + run( + [ + "docker", + "build", + "-t", + image_name, + "-f", + str(dockerfile_path), + str(project_path), + ], + "Error building Docker image locally", + ) + + def list_images(self): + run(["docker", "image", "ls"], "Error listing Docker images locally") + + def remove_image(self, image_name: str): + run( + ["docker", "rmi", "-f", image_name], + f"Error removing Docker image {image_name} locally", + ) + + +class MinikubeDockerHandler(BaseDockerHandler): + def __init__(self, minikube_helper: MinikubeHelper): + self.minikube_helper = minikube_helper + self.env = self._get_minikube_env() + + def is_running(self): + if not self.minikube_helper.is_running(): + raise RuntimeError( + "Minikube not running. Please run the Gaiaflow services in prod-local mode first." + ) + + def build(self, image_name: str, dockerfile_path: Path, project_path: Path): + self.is_running() + log_info(f"Building Docker image [{image_name}] in Minikube context") + run( + [ + "docker", + "build", + "-t", + image_name, + "-f", + str(dockerfile_path), + str(project_path), + ], + "Error building Docker image inside Minikube", + env=self.env, + ) + + def list_images(self): + self.is_running() + run( + ["docker", "image", "ls"], + "Error listing Docker images inside Minikube", + env=self.env, + ) + + def remove_image(self, image_name: str): + self.is_running() + run( + ["docker", "rmi", "-f", image_name], + f"Error removing Docker image {image_name} inside Minikube", + env=self.env, + ) + + def _get_minikube_env(self): + self.is_running() + result = self.minikube_helper.run_cmd( + ["docker-env", "--shell", "bash"], stdout=subprocess.PIPE, check=True + ) + return MinikubeDockerHandler._parse_minikube_env(result.stdout.decode()) + + @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 + + +class LocalUserCustomImageDockerHandler(LocalDockerHandler): + def pre_build(self, image_name: str, dockerfile_path: Path, project_path: Path): + pass + + def build(self, image_name: str, dockerfile_path: Path, project_path: Path): + log_info("Building user provided dockerfile") + super().build(image_name, dockerfile_path, project_path) + + +class MinikubeUserCustomImageDockerHandler(MinikubeDockerHandler): + def pre_build(self, image_name: str, dockerfile_path: Path, project_path: Path): + pass + + def build(self, image_name: str, dockerfile_path: Path, project_path: Path): + log_info("Building user provided dockerfile") + super().build(image_name, dockerfile_path, project_path) + + +HANDLER_REGISTRY = { + DockerHandlerMode.LOCAL: LocalDockerHandler, + DockerHandlerMode.MINIKUBE: MinikubeDockerHandler, + DockerHandlerMode.LOCAL_USER: LocalUserCustomImageDockerHandler, + DockerHandlerMode.MINIKUBE_USER: MinikubeUserCustomImageDockerHandler, +} + + +class DockerHelper: + def __init__(self, image_name: str, project_path: Path, handler: BaseDockerHandler): + self.image_name = image_name + self.project_path = project_path + self.handler = handler + + def build_image(self, dockerfile_path: Path): + if not dockerfile_path.exists(): + log_error(f"Dockerfile not found at {dockerfile_path}") + return + + pre_build_ctx = self.handler.pre_build( + self.image_name, dockerfile_path, self.project_path + ) + if pre_build_ctx: + with pre_build_ctx: + self.handler.build(self.image_name, dockerfile_path, self.project_path) + else: + self.handler.build(self.image_name, dockerfile_path, self.project_path) + + self.handler.post_build(self.image_name, dockerfile_path, self.project_path) + + def list_images(self): + self.handler.list_images() + + def remove_image(self, image_name: str): + self.handler.remove_image(image_name) + + def prune_images(self): + self.handler.prune_images() + + +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): + if self.os_type == "linux" or is_wsl(): + 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 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" / "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", + stdout=f, + ) + log_info(f"Created kube config inline file {filename}") diff --git a/src/gaiaflow/managers/minikube_manager.py b/src/gaiaflow/managers/minikube_manager.py index 7b573f7..82a6bac 100644 --- a/src/gaiaflow/managers/minikube_manager.py +++ b/src/gaiaflow/managers/minikube_manager.py @@ -1,12 +1,7 @@ -import os import platform -import shutil import subprocess -from contextlib import contextmanager from pathlib import Path -from typing import Any, Set, Literal - -import yaml +from typing import Any, Literal, Set from gaiaflow.constants import ( AIRFLOW_SERVICES, @@ -17,357 +12,20 @@ ExtendedAction, ) from gaiaflow.managers.base_manager import BaseGaiaflowManager -from gaiaflow.managers.mlops_manager import MlopsManager -from gaiaflow.managers.utils import ( - find_python_packages, - is_wsl, - log_error, - log_info, - run, - set_permissions, +from gaiaflow.managers.helpers import ( + BaseDockerHandler, + DockerHelper, + KubeConfigHelper, + MinikubeHelper, + DockerHandlerMode, ) - - -@contextmanager -def temporary_copy(src: Path, dest: Path): - print("copying...", src, dest) - shutil.copyfile(src, dest) - try: - yield - finally: - if dest.exists(): - 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 BaseDockerHandler: - """Abstract docker handler with optional hooks.""" - - @classmethod - def get_docker_handler(cls, mode: str, **kwargs): - handler_cls = HANDLER_REGISTRY.get(mode) - if not handler_cls: - raise ValueError(f"Unknown Docker build mode: {mode}") - return handler_cls(**kwargs) - - def pre_build(self, image_name: str, dockerfile_path: Path, project_path: Path): - """Override this if you want a different or no pre_build""" - log_info(f"Updating Dockerfile at {dockerfile_path}") - BaseDockerHandler._add_copy_statements_to_dockerfile( - str(dockerfile_path), find_python_packages(project_path) - ) - runner_src = Path(__file__).parent.parent.resolve() / "core" / "runner.py" - runner_dest = project_path / "runner.py" - return temporary_copy(runner_src, runner_dest) - - def build(self, image_name: str, dockerfile_path: Path, project_path: Path): - raise NotImplementedError - - def post_build(self, image_name: str, dockerfile_path: Path, project_path: Path): - pass - - def list_images(self): - raise NotImplementedError - - def remove_image(self, image_name: str): - raise NotImplementedError - - def prune_images(self): - raise NotImplementedError - - def _update_dockerfile(self, dockerfile_path: Path): - BaseDockerHandler._add_copy_statements_to_dockerfile( - str(dockerfile_path), find_python_packages(self.project_path) - ) - - @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 LocalDockerHandler(BaseDockerHandler): - def build(self, image_name: str, dockerfile_path: Path, project_path: Path): - log_info(f"Building Docker image [{image_name}] locally") - run( - ["docker", "build", "-t", image_name, "-f", str(dockerfile_path), str(project_path)], - "Error building Docker image locally", - ) - - def list_images(self): - run(["docker", "image", "ls"], "Error listing Docker images locally") - - def remove_image(self, image_name: str): - run(["docker", "rmi", "-f", image_name], f"Error removing Docker image {image_name} " - "locally") - - def prune_images(self): - run(["docker", "image", "prune", "-f"], "Error pruning Docker images " - "locally") - -class MinikubeDockerHandler(BaseDockerHandler): - def __init__(self, minikube_helper: MinikubeHelper): - self.minikube_helper = minikube_helper - self.env = self._get_minikube_env() - - def build(self, image_name: str, dockerfile_path: Path, project_path: Path): - log_info(f"Building Docker image [{image_name}] in Minikube context") - run( - ["docker", "build", "-t", image_name, "-f", str(dockerfile_path), str(project_path)], - "Error building Docker image inside Minikube", - env=self.env, - ) - - def list_images(self): - run(["docker", "image", "ls"],"Error listing Docker images inside Minikube", env=self.env) - - def remove_image(self, image_name: str): - run(["docker", "rmi", "-f", image_name], f"Error removing Docker image {image_name} " - "inside Minikube", env=self.env) - - def prune_images(self): - run(["docker", "image", "prune", "-f"], "Error pruning Docker images " - "inside Minikube", env=self.env) - - def _get_minikube_env(self): - result = self.minikube_helper.run_cmd( - ["docker-env", "--shell", "bash"], stdout=subprocess.PIPE, check=True - ) - return MinikubeDockerHandler._parse_minikube_env(result.stdout.decode()) - - @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 - -class LocalUserCustomImageDockerHandler(LocalDockerHandler): - def pre_build(self, image_name: str, dockerfile_path: Path, project_path: Path): - pass - - def build(self, image_name: str, dockerfile_path: Path, project_path: Path): - log_info("Building user provided dockerfile") - super().build(image_name, dockerfile_path, project_path) - -class MinikubeUserCustomImageDockerHandler(MinikubeDockerHandler): - def pre_build(self, image_name: str, dockerfile_path: Path, project_path: Path): - pass - - def build(self, image_name: str, dockerfile_path: Path, project_path: Path): - log_info("Building user provided dockerfile") - super().build(image_name, dockerfile_path, project_path) - -HANDLER_REGISTRY = { - "local": LocalDockerHandler, - "minikube": MinikubeDockerHandler, - "local-user": LocalUserCustomImageDockerHandler, - "minikube-user": MinikubeUserCustomImageDockerHandler -} - -class DockerHelper: - def __init__( - self, - image_name: str, - project_path: Path, - builder: BaseDockerHandler - ): - self.image_name = image_name - self.project_path = project_path - self.builder = builder - - def build_image(self, dockerfile_path: Path): - if not dockerfile_path.exists(): - log_error(f"Dockerfile not found at {dockerfile_path}") - return - - pre_build_ctx = self.builder.pre_build( - self.image_name, dockerfile_path, self.project_path - ) - if pre_build_ctx: - with pre_build_ctx: - self.builder.build(self.image_name, dockerfile_path, self.project_path) - else: - self.builder.build(self.image_name, dockerfile_path, self.project_path) - - self.builder.post_build(self.image_name, dockerfile_path, self.project_path) - - def list_images(self): - self.builder.list_images() - - def remove_image(self, image_name: str): - self.builder.remove_image(image_name) - - def prune_images(self): - self.builder.prune_images() - - -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): - if self.os_type == "linux" or is_wsl(): - 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 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" / "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", - stdout=f, - ) - log_info(f"Created kube config inline file {filename}") +from gaiaflow.managers.mlops_manager import MlopsManager +from gaiaflow.managers.utils import log_info, run class MinikubeManager(BaseGaiaflowManager): allowed_kwargs = {"secret_name", "secret_data", "dockerfile_path"} + def __init__( self, gaiaflow_path: Path, @@ -375,7 +33,7 @@ def __init__( action: Action, force_new: bool = False, prune: bool = False, - docker_build_mode: Literal["local", "minikube"] = "local", + docker_handler_mode: DockerHandlerMode = DockerHandlerMode.LOCAL, image_name: str = "", **kwargs, ): @@ -392,12 +50,13 @@ def __init__( self.image_name = image_name self.minikube_helper = MinikubeHelper() - builder = BaseDockerHandler.get_docker_builder(docker_build_mode, - minikube_helper=self.minikube_helper) + handler = BaseDockerHandler.get_docker_handler( + docker_handler_mode, minikube_helper=self.minikube_helper + ) self.docker_helper = DockerHelper( image_name=image_name, project_path=user_project_path, - builder=builder, + handler=handler, ) self.kube_helper = KubeConfigHelper( gaiaflow_path=gaiaflow_path, os_type=self.os_type @@ -416,6 +75,8 @@ def _get_valid_actions(self) -> Set[Action]: ExtendedAction.DOCKERIZE, ExtendedAction.CREATE_CONFIG, ExtendedAction.CREATE_SECRET, + ExtendedAction.LIST_IMAGES, + ExtendedAction.REMOVE_IMAGE, } @classmethod @@ -432,11 +93,15 @@ def run(cls, **kwargs): BaseAction.RESTART: manager.restart, BaseAction.CLEANUP: manager.cleanup, ExtendedAction.DOCKERIZE: lambda: manager.build_docker_image( - kwargs["dockerfile_path"]), + kwargs["dockerfile_path"] + ), ExtendedAction.CREATE_CONFIG: manager.create_kube_config_inline, ExtendedAction.CREATE_SECRET: lambda: manager.create_secrets( kwargs["secret_name"], kwargs["secret_data"] ), + ExtendedAction.LIST_IMAGES: manager.list_images, + ExtendedAction.REMOVE_IMAGE: lambda: manager.remove_image( + kwargs["image_name"]), } try: @@ -476,7 +141,9 @@ def create_kube_config_inline(self): def build_docker_image(self, dockerfile_path: str): if not dockerfile_path: - dockerfile_path = self.gaiaflow_path / "_docker" / "user-package" / "Dockerfile" + dockerfile_path = ( + self.gaiaflow_path / "_docker" / "user-package" / "Dockerfile" + ) self.docker_helper.build_image(dockerfile_path) def create_secrets(self, secret_name: str, secret_data: dict[str, Any]): @@ -529,3 +196,12 @@ def cleanup(self): "Error removing airflow docker network", ) log_info("Minikube Cleanup complete") + + def list_images(self): + self.docker_helper.list_images() + + def remove_image(self, image_name: str): + self.docker_helper.remove_image(image_name) + + def prune_images(self): + self.docker_helper.prune_images() diff --git a/src/gaiaflow/managers/mlops_manager.py b/src/gaiaflow/managers/mlops_manager.py index 3b8eb16..59ab2f5 100644 --- a/src/gaiaflow/managers/mlops_manager.py +++ b/src/gaiaflow/managers/mlops_manager.py @@ -2,37 +2,32 @@ import os import platform import shutil -import socket -import subprocess from pathlib import Path from typing import Set import fsspec -import psutil -import yaml from ruamel.yaml import YAML from gaiaflow.constants import ( - AIRFLOW_SERVICES, GAIAFLOW_STATE_FILE, - MINIO_SERVICES, - MLFLOW_SERVICES, Action, BaseAction, ExtendedAction, Service, ) from gaiaflow.managers.base_manager import BaseGaiaflowManager +from gaiaflow.managers.helpers import ( + DockerComposeHelper, + DockerResources, + JupyterHelper, +) from gaiaflow.managers.utils import ( convert_crlf_to_lf, create_directory, delete_project_state, - env_exists, gaiaflow_path_exists_in_state, - handle_error, log_error, log_info, - run, save_project_state, set_permissions, update_entrypoint_install_path, @@ -40,155 +35,6 @@ ) -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", - "alpine/socat", - ] - - AIRFLOW_CONTAINERS = [ - "airflow-apiserver", - "airflow-scheduler", - "airflow-dag-processor", - "airflow-triggerer", - "docker-proxy" - ] - - VOLUMES = [ - "docker-compose_postgres-db-volume-airflow", - "docker-compose_postgres-db-volume-mlflow", - ] - - SERVICES = { - "airflow": AIRFLOW_SERVICES, - "mlflow": MLFLOW_SERVICES, - "minio": MINIO_SERVICES, - } - - -class DockerComposeHelper: - 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/docker-compose/docker-compose.yml", - ] - if self.is_prod_local: - base += [ - "-f", - f"{self.gaiaflow_path}/_docker/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): """Manager class to Start/Stop/Restart MLOps Docker services.""" diff --git a/src/gaiaflow/managers/utils.py b/src/gaiaflow/managers/utils.py index ec7cdd6..dfa4480 100644 --- a/src/gaiaflow/managers/utils.py +++ b/src/gaiaflow/managers/utils.py @@ -98,7 +98,6 @@ def save_project_state(project_path: Path, gaiaflow_path: Path): def load_project_state() -> dict | None: state_file = get_state_file() - print("state_file", state_file) if not state_file.exists(): return None From a2b28d1c73da35387cbc712f798d78c9214e35a4 Mon Sep 17 00:00:00 2001 From: Yogesh Kumar Baljeet Singh Date: Tue, 16 Sep 2025 16:03:11 +0200 Subject: [PATCH 13/15] add and fix tests --- src/gaiaflow/cli/commands/minikube.py | 6 +- src/gaiaflow/managers/helpers.py | 23 +- src/gaiaflow/managers/minikube_manager.py | 9 +- tests/cli/commands/test_minikube.py | 69 ++- tests/cli/commands/test_mlops.py | 57 ++- tests/core/test_operators.py | 4 +- tests/managers/test_helpers.py | 559 ++++++++++++++++++++++ tests/managers/test_minikube_manager.py | 362 ++------------ tests/managers/test_mlops_manager.py | 59 +-- 9 files changed, 713 insertions(+), 435 deletions(-) create mode 100644 tests/managers/test_helpers.py diff --git a/src/gaiaflow/cli/commands/minikube.py b/src/gaiaflow/cli/commands/minikube.py index 33449d8..140cbcd 100644 --- a/src/gaiaflow/cli/commands/minikube.py +++ b/src/gaiaflow/cli/commands/minikube.py @@ -124,14 +124,14 @@ def dockerize( typer.echo("Please create a project with Gaiaflow before running this command.") return if dockerfile_path: - docker_build_mode = "minikube-user" + docker_handler_mode = "minikube-user" else: - docker_build_mode = "minikube" + docker_handler_mode = "minikube" imports.MinikubeManager.run( gaiaflow_path=gaiaflow_path, user_project_path=user_project_path, action=imports.ExtendedAction.DOCKERIZE, - docker_build_mode=docker_build_mode, + docker_handler_mode=docker_handler_mode, image_name=image_name, dockerfile_path=dockerfile_path, ) diff --git a/src/gaiaflow/managers/helpers.py b/src/gaiaflow/managers/helpers.py index 3fffb57..2ecb400 100644 --- a/src/gaiaflow/managers/helpers.py +++ b/src/gaiaflow/managers/helpers.py @@ -476,18 +476,17 @@ def __init__(self, gaiaflow_path: Path, os_type: str): self.os_type = os_type def create_inline(self): - if self.os_type == "linux" or is_wsl(): - 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 backup_config.exists(): - shutil.copy(backup_config, kube_config) - backup_config.unlink() - log_info("Reverted kube config to original state.") + 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 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(): diff --git a/src/gaiaflow/managers/minikube_manager.py b/src/gaiaflow/managers/minikube_manager.py index 82a6bac..cfbbfcf 100644 --- a/src/gaiaflow/managers/minikube_manager.py +++ b/src/gaiaflow/managers/minikube_manager.py @@ -139,8 +139,8 @@ def stop(self): def create_kube_config_inline(self): self.kube_helper.create_inline() - def build_docker_image(self, dockerfile_path: str): - if not dockerfile_path: + def build_docker_image(self, dockerfile_path: str = ""): + if dockerfile_path != "": dockerfile_path = ( self.gaiaflow_path / "_docker" / "user-package" / "Dockerfile" ) @@ -201,7 +201,4 @@ def list_images(self): self.docker_helper.list_images() def remove_image(self, image_name: str): - self.docker_helper.remove_image(image_name) - - def prune_images(self): - self.docker_helper.prune_images() + self.docker_helper.remove_image(image_name) \ No newline at end of file diff --git a/tests/cli/commands/test_minikube.py b/tests/cli/commands/test_minikube.py index 269bf76..bf2341c 100644 --- a/tests/cli/commands/test_minikube.py +++ b/tests/cli/commands/test_minikube.py @@ -8,6 +8,8 @@ ExtendedAction from gaiaflow.cli.commands.minikube import app as prod_app from gaiaflow.cli.commands.minikube import load_imports +from gaiaflow.managers.helpers import DockerHandlerMode + class TestGaiaflowProdCLI(unittest.TestCase): def setUp(self): @@ -32,15 +34,11 @@ def test_start_command(self, mock_load_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 ) @@ -59,7 +57,6 @@ def test_start_command_exits_when_project_not_exists(self, mock_load_imports): result = self.runner.invoke(prod_app, [ "start", - "--path", str(self.test_project_path) ]) self.assertEqual(result.exit_code, 0) @@ -73,7 +70,6 @@ def test_stop_command(self, mock_load_imports): result = self.runner.invoke(prod_app, [ "stop", - "--path", str(self.test_project_path) ]) self.assertEqual(result.exit_code, 0) @@ -90,7 +86,6 @@ def test_restart_command(self, mock_load_imports): result = self.runner.invoke(prod_app, [ "restart", - "--path", str(self.test_project_path), "--force-new" ]) @@ -100,6 +95,7 @@ def test_restart_command(self, mock_load_imports): gaiaflow_path=self.test_gaiaflow_path, user_project_path=self.test_project_path, action=BaseAction.RESTART, + force_new=True ) @patch('gaiaflow.cli.commands.minikube.load_imports') @@ -108,7 +104,6 @@ def test_dockerize_command(self, mock_load_imports): result = self.runner.invoke(prod_app, [ "dockerize", - "--path", str(self.test_project_path), "--image-name", "my-custom-image" ]) @@ -118,8 +113,42 @@ def test_dockerize_command(self, mock_load_imports): gaiaflow_path=self.test_gaiaflow_path, user_project_path=self.test_project_path, action=ExtendedAction.DOCKERIZE, - local=False, - image_name="my-custom-image" + docker_handler_mode=DockerHandlerMode.MINIKUBE, + image_name="my-custom-image", dockerfile_path=None + ) + + @patch("gaiaflow.cli.commands.minikube.load_imports") + def test_list_images_command(self, mock_load_imports): + mock_load_imports.return_value = self.mock_imports + + result = self.runner.invoke( + prod_app, ["list-images"] + ) + + 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.LIST_IMAGES, + docker_handler_mode=DockerHandlerMode.MINIKUBE, + ) + + @patch("gaiaflow.cli.commands.minikube.load_imports") + def test_remove_image_command(self, mock_load_imports): + mock_load_imports.return_value = self.mock_imports + + result = self.runner.invoke(prod_app, ["remove-image", "--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.REMOVE_IMAGE, + docker_handler_mode=DockerHandlerMode.MINIKUBE, + image_name="my-custom-image", ) @patch('gaiaflow.cli.commands.minikube.load_imports') @@ -128,7 +157,6 @@ def test_dockerize_command_with_default_image_name(self, mock_load_imports): result = self.runner.invoke(prod_app, [ "dockerize", - "--path", str(self.test_project_path) ]) self.assertEqual(result.exit_code, 0) @@ -143,7 +171,6 @@ def test_create_config_command(self, mock_load_imports): result = self.runner.invoke(prod_app, [ "create-config", - "--path", str(self.test_project_path) ]) self.assertEqual(result.exit_code, 0) @@ -160,7 +187,6 @@ def test_create_secret_command(self, mock_load_imports): result = self.runner.invoke(prod_app, [ "create-secret", - "--path", str(self.test_project_path), "--name", "my-secret", "--data", "key1=value1", "--data", "key2=value2" @@ -186,7 +212,6 @@ def test_cleanup_command(self, mock_load_imports): result = self.runner.invoke(prod_app, [ "cleanup", - "--path", str(self.test_project_path) ]) self.assertEqual(result.exit_code, 0) @@ -203,14 +228,13 @@ def test_all_commands_handle_missing_project_gracefully(self, mock_load_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)], + ["start"], + ["stop"], + ["restart"], + ["dockerize"], + ["create-config"], + ["create-secret", "--name", "test", "--data", "key=value"], + ["cleanup"], ] for command_args in commands_and_args: @@ -244,7 +268,6 @@ def test_action_objects_comparison(self, mock_load_imports): result = self.runner.invoke(prod_app, [ "start", - "--path", str(self.test_project_path) ]) self.assertEqual(result.exit_code, 0) diff --git a/tests/cli/commands/test_mlops.py b/tests/cli/commands/test_mlops.py index b95126f..124e0ff 100644 --- a/tests/cli/commands/test_mlops.py +++ b/tests/cli/commands/test_mlops.py @@ -7,6 +7,8 @@ from gaiaflow.constants import Service, DEFAULT_IMAGE_NAME, BaseAction, \ ExtendedAction from gaiaflow.cli.commands.mlops import app, load_imports +from gaiaflow.managers.helpers import DockerHandlerMode + class TestGaiaflowCLI(unittest.TestCase): def setUp(self): @@ -32,14 +34,10 @@ def test_start_command_with_all_services(self, mock_load_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 ) @@ -63,7 +61,6 @@ def test_start_command_with_specific_services(self, mock_load_imports): result = self.runner.invoke(app, [ "start", - "--path", str(self.test_project_path), "--service", "jupyter", "--service", "airflow", "--cache", @@ -112,7 +109,6 @@ def test_start_command_saves_project_state_when_not_exists(self, mock_load_impor result = self.runner.invoke(app, [ "start", - "--path", str(self.test_project_path) ]) self.assertEqual(result.exit_code, 0) @@ -126,7 +122,6 @@ def test_stop_command_with_all_services(self, mock_load_imports): result = self.runner.invoke(app, [ "stop", - "--path", str(self.test_project_path), "--delete-volume" ]) @@ -146,7 +141,6 @@ def test_stop_command_with_specific_services(self, mock_load_imports): result = self.runner.invoke(app, [ "stop", - "--path", str(self.test_project_path), "--service", "jupyter" ]) @@ -166,7 +160,6 @@ def test_restart_command_with_all_services(self, mock_load_imports): result = self.runner.invoke(app, [ "restart", - "--path", str(self.test_project_path), "--force-new", "--cache", "--jupyter-port", "9001", @@ -196,8 +189,6 @@ def test_restart_command_with_specific_services(self, mock_load_imports): app, [ "restart", - "--path", - str(self.test_project_path), "--service", "jupyter", "--service", @@ -241,7 +232,6 @@ def test_cleanup_command(self, mock_load_imports): result = self.runner.invoke(app, [ "cleanup", - "--path", str(self.test_project_path), "--prune" ]) @@ -261,7 +251,6 @@ def test_dockerize_command(self, mock_load_imports): result = self.runner.invoke(app, [ "dockerize", - "--path", str(self.test_project_path), "--image-name", "custom-image" ]) @@ -275,17 +264,50 @@ def test_dockerize_command(self, mock_load_imports): gaiaflow_path=self.test_gaiaflow_path, user_project_path=self.test_project_path, action=ExtendedAction.DOCKERIZE, - local=True, + docker_build_mode=DockerHandlerMode.LOCAL, image_name="custom-image" ) + @patch("gaiaflow.cli.commands.mlops.load_imports") + def test_list_images_command(self, mock_load_imports): + mock_load_imports.return_value = self.mock_imports + + result = self.runner.invoke(app, ["list-images"]) + + 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.LIST_IMAGES, + docker_handler_mode=DockerHandlerMode.LOCAL, + ) + + @patch("gaiaflow.cli.commands.mlops.load_imports") + def test_remove_image_command(self, mock_load_imports): + mock_load_imports.return_value = self.mock_imports + + result = self.runner.invoke( + app, ["remove-image", "--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.REMOVE_IMAGE, + docker_handler_mode=DockerHandlerMode.LOCAL, + image_name="my-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) @@ -299,7 +321,7 @@ 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)] + app, ["update-deps"] ) self.assertEqual(result.exit_code, 0) @@ -321,7 +343,6 @@ def test_commands_handle_missing_project_gracefully(self, mock_load_imports): with self.subTest(command=command): result = self.runner.invoke(app, [ command, - "--path", str(self.test_project_path) ]) self.assertEqual(result.exit_code, 0) @@ -347,14 +368,12 @@ def test_argument_type_conversion(self, mock_load_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__': diff --git a/tests/core/test_operators.py b/tests/core/test_operators.py index 12f9414..daf8a37 100644 --- a/tests/core/test_operators.py +++ b/tests/core/test_operators.py @@ -335,7 +335,7 @@ def test_create_dev_docker_task(self, mock_ext_op): 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["docker_url"], "tcp://docker-proxy:2375") self.assertEqual(kwargs["retrieve_output"], True) self.assertEqual(kwargs["retrieve_output_path"], "/tmp/script.out") self.assertEqual( @@ -376,7 +376,7 @@ def test_create_dev_docker_task_with_custom_env_vars(self, mock_ext_op): 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["docker_url"], "tcp://docker-proxy:2375") self.assertEqual(kwargs["retrieve_output"], True) self.assertEqual(kwargs["retrieve_output_path"], "/tmp/script.out") self.assertEqual( diff --git a/tests/managers/test_helpers.py b/tests/managers/test_helpers.py new file mode 100644 index 0000000..62d9947 --- /dev/null +++ b/tests/managers/test_helpers.py @@ -0,0 +1,559 @@ +import unittest +import tempfile +from pathlib import Path +import socket +import subprocess +from unittest.mock import patch, MagicMock + +import yaml + +from gaiaflow.constants import BaseAction, Service +from gaiaflow.managers.helpers import ( + JupyterHelper, + KubeConfigHelper, + DockerHelper, + MinikubeHelper, + BaseDockerHandler, + MinikubeDockerHandler, + DockerHandlerMode, + temporary_copy, + DockerComposeHelper, + DockerResources, + MinikubeUserCustomImageDockerHandler, + LocalUserCustomImageDockerHandler, + LocalDockerHandler, +) +from gaiaflow.managers.mlops_manager import MlopsManager + + +class TestJupyterHelper(unittest.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" / "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" / "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_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.helpers.env_exists", return_value=True): + helper.start() + mock_popen.assert_called() + + @patch("psutil.process_iter") + def test_stop_terminates_jupyter_processes(self, mock_iter): + fake_proc = MagicMock() + helper = JupyterHelper(8895, "mamba", "custom-env", self.manager.gaiaflow_path) + fake_proc.info = {"pid": 1, "name": "jupyter", "cmdline": ["jupyter-lab"]} + mock_iter.return_value = [fake_proc] + helper.stop() + fake_proc.terminate.assert_called_once() + + +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.helpers.run") + @patch("gaiaflow.managers.helpers.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.helpers.run") + @patch("gaiaflow.managers.helpers.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.helpers.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.helpers.run", + side_effect=subprocess.CalledProcessError(1, "cmd"), + ) + @patch("gaiaflow.managers.helpers.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.helpers.run") + def test_stop(self, mock_run): + self.helper.stop() + mock_run.assert_called() + + @patch("gaiaflow.managers.helpers.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 TestDockerComposeHelper(unittest.TestCase): + def setUp(self): + self.helper = DockerComposeHelper(Path("/fake/path"), is_prod_local=True) + + + def test_base_cmd_with_prod_local(self): + cmd = self.helper._base_cmd() + self.assertIn("docker", cmd) + self.assertIn("docker-compose-minikube-network.yml", " ".join(cmd)) + + + def test_docker_services_for_valid_key(self): + services = DockerComposeHelper.docker_services_for("airflow") + self.assertEqual(services, DockerResources.SERVICES["airflow"]) + + + def test_docker_services_for_invalid_key(self): + services = DockerComposeHelper.docker_services_for("unknown") + self.assertEqual(services, []) + + + @patch("gaiaflow.managers.helpers.run") + @patch("gaiaflow.managers.helpers.log_info") + def test_run_compose_with_service(self, mock_log, mock_run): + with patch.object(DockerComposeHelper, "docker_services_for", return_value=["svc"]): + self.helper.run_compose(["up"], service="airflow") + mock_run.assert_called() + + + @patch("gaiaflow.managers.helpers.run") + def test_prune_runs_expected_cmds(self, mock_run): + DockerComposeHelper.prune() + self.assertGreaterEqual(mock_run.call_count, len(DockerResources.IMAGES)) + +class TestKubeConfigHelper(unittest.TestCase): + def setUp(self): + self.tmpdir = tempfile.TemporaryDirectory() + self.gaia_path = Path(self.tmpdir.name) + (self.gaia_path / "_docker").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" / "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.helpers.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" / "kube_config_inline").exists() + ) + + @patch("gaiaflow.managers.helpers.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.helpers.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.helpers.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 TestBaseDockerHandler(unittest.TestCase): + def setUp(self): + self.tmpdir = tempfile.TemporaryDirectory() + self.project_path = Path(self.tmpdir.name) + self.dockerfile = self.project_path / "Dockerfile" + self.dockerfile.write_text("ENV TEST=1\nENTRYPOINT test.sh\n") + + def tearDown(self): + self.tmpdir.cleanup() + + def test_get_docker_handler_valid_and_invalid(self): + handler = BaseDockerHandler.get_docker_handler(DockerHandlerMode.LOCAL) + self.assertIsInstance(handler, LocalDockerHandler) + with self.assertRaises(ValueError): + BaseDockerHandler.get_docker_handler("bad-mode") + + @patch("gaiaflow.managers.helpers.find_python_packages", return_value=[ + "pkg1"]) + @patch("gaiaflow.managers.helpers.temporary_copy") + def test_pre_build_returns_contextmanager(self, mock_temp_copy, _): + handler = BaseDockerHandler() + ctx = handler.pre_build("img", self.dockerfile, self.project_path) + self.assertTrue(hasattr(ctx, "__enter__")) + mock_temp_copy.assert_called_once() + + def test_add_copy_statements_writes_expected_lines(self): + BaseDockerHandler._add_copy_statements_to_dockerfile(str(self.dockerfile), ["mypkg"]) + text = self.dockerfile.read_text() + self.assertIn("COPY mypkg ./mypkg", text) + self.assertIn("COPY runner.py ./runner.py", text) + + @patch("gaiaflow.managers.helpers.find_python_packages", return_value=["mypkg"]) + def test_base_dockerhandler_update_dockerfile(self, _): + handler = BaseDockerHandler() + handler.project_path = ( + self.project_path + ) + handler._update_dockerfile(self.dockerfile) + + contents = self.dockerfile.read_text() + self.assertIn("COPY mypkg ./mypkg", contents) + self.assertIn("COPY runner.py ./runner.py", contents) + env_index = contents.splitlines().index("ENV TEST=1") + entry_index = contents.splitlines().index("ENTRYPOINT test.sh") + copy_index = contents.splitlines().index("COPY mypkg ./mypkg") + self.assertGreater(copy_index, env_index) + self.assertLess(copy_index, entry_index) + + def test_abstract_methods_raise(self): + handler = BaseDockerHandler() + with self.assertRaises(NotImplementedError): + handler.build("img", self.dockerfile, self.project_path) + with self.assertRaises(NotImplementedError): + handler.list_images() + with self.assertRaises(NotImplementedError): + handler.remove_image("img") + + +class TestDockerHandlers(unittest.TestCase): + def setUp(self): + self.tmpdir = tempfile.TemporaryDirectory() + self.project_path = Path(self.tmpdir.name) + self.dockerfile = self.project_path / "Dockerfile" + self.dockerfile.write_text("ENV TEST=1\nENTRYPOINT test.sh\n") + + + def tearDown(self): + self.tmpdir.cleanup() + + @patch("gaiaflow.managers.helpers.run") + def test_local_docker_handler_build(self, mock_run): + handler = LocalDockerHandler() + handler.build("img", self.dockerfile, self.project_path) + + mock_run.assert_called_once() + cmd, msg = mock_run.call_args[0] + self.assertEqual( + cmd, + [ + "docker", + "build", + "-t", + "img", + "-f", + str(self.dockerfile), + str(self.project_path), + ], + ) + self.assertIn("Error building Docker image locally", msg) + + + @patch("gaiaflow.managers.helpers.run") + def test_local_docker_handler_list(self, mock_run): + handler = LocalDockerHandler() + handler.list_images() + handler.remove_image("img") + + cmd1, msg1 = mock_run.call_args_list[0][0] + self.assertEqual(cmd1, ["docker", "image", "ls"]) + self.assertIn("Error listing Docker images locally", msg1) + + @patch("gaiaflow.managers.helpers.run") + def test_local_docker_handler_remove(self, mock_run): + handler = LocalDockerHandler() + handler.list_images() + handler.remove_image("img") + + cmd2, msg2 = mock_run.call_args_list[1][0] + self.assertEqual(cmd2, ["docker", "rmi", "-f", "img"]) + self.assertIn("Error removing Docker image img locally", msg2) + + + @patch.object(MinikubeHelper, "is_running", return_value=True) + @patch("gaiaflow.managers.helpers.run") + def test_minikube_docker_handler_build(self, mock_run, _): + handler = MinikubeDockerHandler(MinikubeHelper()) + handler.build("img", self.dockerfile, self.project_path) + + mock_run.assert_called_once() + cmd, msg = mock_run.call_args[0] + self.assertEqual( + cmd, + [ + "docker", + "build", + "-t", + "img", + "-f", + str(self.dockerfile), + str(self.project_path), + ], + ) + self.assertIn("Error building Docker image inside Minikube", msg) + self.assertIn("env", mock_run.call_args.kwargs) + + @patch.object(MinikubeHelper, "is_running", return_value=True) + @patch("gaiaflow.managers.helpers.run") + def test_minikube_docker_handler_list(self, mock_run, _): + handler = MinikubeDockerHandler(MinikubeHelper()) + handler.list_images() + handler.remove_image("img") + + cmd1, msg1 = mock_run.call_args_list[0][0] + self.assertEqual(cmd1, ["docker", "image", "ls"]) + self.assertIn("Error listing Docker images inside Minikube", msg1) + self.assertIn("env", mock_run.call_args_list[0].kwargs) + + @patch.object(MinikubeHelper, "is_running", return_value=True) + @patch("gaiaflow.managers.helpers.run") + def test_minikube_docker_handler_remove(self, mock_run, _): + handler = MinikubeDockerHandler(MinikubeHelper()) + handler.list_images() + handler.remove_image("img") + + cmd2, msg2 = mock_run.call_args_list[1][0] + self.assertEqual(cmd2, ["docker", "rmi", "-f", "img"]) + self.assertIn("Error removing Docker image img inside Minikube", msg2) + self.assertIn("env", mock_run.call_args_list[1].kwargs) + + + def test_parse_minikube_env(self): + output = 'export FOO="bar"\nexport BAZ="qux"\n' + env = MinikubeDockerHandler._parse_minikube_env(output) + self.assertEqual(env["FOO"], "bar") + self.assertEqual(env["BAZ"], "qux") + + + def test_local_user_custom_handler_build(self): + handler = LocalUserCustomImageDockerHandler() + with patch.object(LocalDockerHandler, "build") as mock_build: + handler.build("img", self.dockerfile, self.project_path) + mock_build.assert_called_once_with("img", self.dockerfile, self.project_path) + + def test_local_user_custom_handler_pre_build_is_none(self): + handler = LocalUserCustomImageDockerHandler() + result = handler.pre_build("img", self.dockerfile, self.project_path) + self.assertIsNone(result) + + def test_minikube_user_custom_handler_build(self): + with patch.object(MinikubeDockerHandler, "build") as mock_build: + handler = MinikubeUserCustomImageDockerHandler(MinikubeHelper()) + handler.build("img", self.dockerfile, self.project_path) + mock_build.assert_called_once_with("img", self.dockerfile, self.project_path) + + def test_minikube_user_custom_handler_pre_build_is_none(self): + handler = MinikubeUserCustomImageDockerHandler(MinikubeHelper()) + result = handler.pre_build("img", self.dockerfile, self.project_path) + self.assertIsNone(result) + + +class TestDockerHelperClass(unittest.TestCase): + def setUp(self): + self.tmpdir = tempfile.TemporaryDirectory() + self.project_path = Path(self.tmpdir.name) + self.dockerfile = self.project_path / "Dockerfile" + self.dockerfile.write_text("FROM alpine\n") + self.handler = MagicMock(spec=BaseDockerHandler) + self.helper = DockerHelper("test-img", self.project_path, self.handler) + + def tearDown(self): + self.tmpdir.cleanup() + + @patch("gaiaflow.managers.helpers.log_error") + def test_build_image_missing_dockerfile_logs_error(self, mock_log): + missing = self.project_path / "DoesNotExist" + self.helper.build_image(missing) + mock_log.assert_called_once_with(f"Dockerfile not found at {missing}") + self.handler.build.assert_not_called() + self.handler.post_build.assert_not_called() + + def test_build_image_with_prebuild_context(self): + cm = MagicMock() + cm.__enter__ = MagicMock() + cm.__exit__ = MagicMock() + self.handler.pre_build.return_value = cm + + self.helper.build_image(self.dockerfile) + + self.handler.pre_build.assert_called_once_with("test-img", self.dockerfile, self.project_path) + cm.__enter__.assert_called_once() + self.handler.build.assert_called_once_with("test-img", self.dockerfile, self.project_path) + self.handler.post_build.assert_called_once_with("test-img", self.dockerfile, self.project_path) + + def test_build_image_without_prebuild_context(self): + self.handler.pre_build.return_value = None + + self.helper.build_image(self.dockerfile) + + self.handler.build.assert_called_once_with("test-img", self.dockerfile, self.project_path) + self.handler.post_build.assert_called_once_with("test-img", self.dockerfile, self.project_path) + + def test_list_images_delegates(self): + self.helper.list_images() + self.handler.list_images.assert_called_once() + + def test_remove_image_delegates(self): + self.helper.remove_image("some-img") + self.handler.remove_image.assert_called_once_with("some-img") diff --git a/tests/managers/test_minikube_manager.py b/tests/managers/test_minikube_manager.py index 05f39f5..dc3a9bd 100644 --- a/tests/managers/test_minikube_manager.py +++ b/tests/managers/test_minikube_manager.py @@ -7,341 +7,14 @@ import yaml from gaiaflow.constants import BaseAction, ExtendedAction +from gaiaflow.managers.helpers import temporary_copy, DockerHandlerMode 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").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" / "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" / "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() @@ -349,7 +22,7 @@ def setUp(self): gaiaflow_path=Path(self.tmpdir.name), user_project_path=Path(self.tmpdir.name), action=BaseAction.START, - local=True, + docker_handler_mode=DockerHandlerMode.MINIKUBE, image_name="img", ) @@ -357,6 +30,26 @@ def setUp(self): def tearDown(self): self.tmpdir.cleanup() + def test_init_with_allowed_kwargs(self): + mgr = MinikubeManager( + Path(self.tmpdir.name), + Path(self.tmpdir.name), + BaseAction.START, + secret_name="my-secret", + secret_data={"k": "v"}, + dockerfile_path=Path("/tmp/Dockerfile"), + ) + self.assertIsInstance(mgr, MinikubeManager) + + def test_init_with_unexpected_kwarg_raises(self): + with self.assertRaises(TypeError) as ctx: + MinikubeManager( + Path(self.tmpdir.name), + Path(self.tmpdir.name), + BaseAction.START, + bad_arg="oops", + ) + def test_valid_actions(self): actions = self.manager._get_valid_actions() @@ -433,3 +126,14 @@ def test_create_kube_config_inline(self, mock_inline): def test_build_docker_image(self, mock_build): self.manager.build_docker_image() mock_build.assert_called() + + @patch.object(DockerHelper, "list_images") + def test_list_docker_images(self, mock_list_images): + self.manager.list_images() + mock_list_images.assert_called() + + @patch.object(DockerHelper, "remove_image") + def test_remove_docker_images(self, mock_remove_image): + self.manager.remove_image("img") + mock_remove_image.assert_called_with("img") + diff --git a/tests/managers/test_mlops_manager.py b/tests/managers/test_mlops_manager.py index b1455bb..4cd6240 100644 --- a/tests/managers/test_mlops_manager.py +++ b/tests/managers/test_mlops_manager.py @@ -56,9 +56,9 @@ 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") + @patch("gaiaflow.managers.helpers.run") + @patch("gaiaflow.managers.helpers.subprocess.Popen") + @patch("gaiaflow.managers.helpers.env_exists") def test_run_dispatches_start(self, mock_env_exists, mock_popen, mock_run): mock_env_exists.return_value = True @@ -87,9 +87,9 @@ def test_run_dispatches_start(self, mock_env_exists, mock_popen, 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") + @patch("gaiaflow.managers.helpers.run") + @patch("gaiaflow.managers.helpers.subprocess.Popen") + @patch("gaiaflow.managers.helpers.env_exists") def test_run_dispatches_start_jupyter_custom_values(self, mock_env_exists, mock_popen, mock_run): @@ -129,16 +129,16 @@ def test_run_invalid_action(self): action="not-an-action", ) - @patch("gaiaflow.managers.mlops_manager.env_exists") - @patch("gaiaflow.managers.mlops_manager.subprocess.Popen") + @patch("gaiaflow.managers.helpers.env_exists") + @patch("gaiaflow.managers.helpers.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") + @patch("gaiaflow.managers.helpers.env_exists") + @patch("gaiaflow.managers.helpers.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, \ @@ -167,7 +167,7 @@ def test_run_missing_action_raises(self): ) self.assertIn("Missing required argument 'action'", str(ctx.exception)) - @patch("gaiaflow.managers.mlops_manager.run") + @patch("gaiaflow.managers.helpers.run") @patch.object(MlopsManager, "_build_docker_images") def test_start_triggers_build_docker_images(self, mock_build, mock_run): self.manager.docker_build = True @@ -261,13 +261,13 @@ def test_update_env_file_updates_existing(self): 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"): + with patch("gaiaflow.managers.mlops_manager.set_permissions"): self.manager._update_files() compose_path = self.gaiaflow_context / "_docker" / "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)) + print(vols) + self.assertTrue(any("project" 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") @@ -275,29 +275,6 @@ 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() @@ -307,7 +284,7 @@ def test_stop_jupyter_processes(self, mock_iter): 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("gaiaflow.managers.helpers.env_exists", return_value=False) @patch("subprocess.Popen") def test_start_jupyter_env_not_exists(self, mock_popen, mock_env): env_name = "test-env" @@ -335,8 +312,8 @@ def test_docker_services_for_known_and_unknown(self): 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") + @patch("gaiaflow.managers.helpers.handle_error") + @patch("gaiaflow.managers.helpers.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=[]): @@ -345,6 +322,6 @@ def test_run_compose_with_service_and_unknown_service(self, mock_run, mock_handl def test_docker_prune(self): helper = self.manager.docker - with patch("gaiaflow.managers.mlops_manager.run") as mock_run: + with patch("gaiaflow.managers.helpers.run") as mock_run: helper.prune() self.assertGreaterEqual(mock_run.call_count, len(DockerResources.IMAGES)) From 793b299b2af54005d6deba6000eaa198bd1cc0f5 Mon Sep 17 00:00:00 2001 From: Yogesh Kumar Baljeet Singh Date: Tue, 16 Sep 2025 16:47:31 +0200 Subject: [PATCH 14/15] improve kube config generation and fix its tests --- pyproject.toml | 1 + src/gaiaflow/managers/helpers.py | 41 +++++------- tests/managers/test_helpers.py | 108 ++++++++++++------------------- 3 files changed, 58 insertions(+), 92 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 1c35aa6..10ecbf8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -56,6 +56,7 @@ ruff = "*" [tool.pixi.tasks] test = "pytest tests" +cov = "pytest --cov=gaiaflow --cov-report=html" format = "isort src/gaiaflow && ruff format src/gaiaflow" check = "ruff check src/gaiaflow" check_fix = "ruff check src/gaiaflow --fix" diff --git a/src/gaiaflow/managers/helpers.py b/src/gaiaflow/managers/helpers.py index 2ecb400..67cf0a6 100644 --- a/src/gaiaflow/managers/helpers.py +++ b/src/gaiaflow/managers/helpers.py @@ -477,31 +477,24 @@ def __init__(self, gaiaflow_path: Path, os_type: str): 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 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(): + log_info("No kube config found, skipping inline creation.") return - with open(kube_config, "r") as f: - config_data = yaml.safe_load(f) + filename = self.gaiaflow_path / "_docker" / "kube_config_inline" + self._write_inline(filename) + + with open(filename, "r") as f: + try: + config_data = yaml.safe_load(f) or {} + except yaml.YAMLError as e: + log_error(f"Failed to parse kube config: {e}") + return + + self._patch_kube_config(config_data, filename) + def _patch_kube_config(self, config_data: dict, filename: Path) -> dict: + log_info("Patching kube config file...") for cluster in config_data.get("clusters", []): cluster_info = cluster.get("cluster", {}) if self.os_type == "windows": @@ -515,11 +508,11 @@ def _patch_kube_config(self, kube_config: Path): cluster_info["server"] = "https://192.168.49.2:8443" cluster_info["insecure-skip-tls-verify"] = True - with open(kube_config, "w") as f: + with open(filename, "w") as f: yaml.dump(config_data, f) + log_info(f"Patched file written to {filename}") - def _write_inline(self, kube_config: Path): - filename = self.gaiaflow_path / "_docker" / "kube_config_inline" + def _write_inline(self, filename: Path): log_info("Creating kube config inline file...") with open(filename, "w") as f: subprocess.call( diff --git a/tests/managers/test_helpers.py b/tests/managers/test_helpers.py index 62d9947..a12b704 100644 --- a/tests/managers/test_helpers.py +++ b/tests/managers/test_helpers.py @@ -236,88 +236,60 @@ 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 + inline_file = self.gaia_path / "_docker" / "kube_config_inline" + inline_file.parent.mkdir(parents=True, exist_ok=True) + with open(inline_file, "w") as f: + yaml.safe_dump(data, f) + return inline_file @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" / "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()) + def test_write_inline_creates_file(self, mock_call): + inline_file = self.gaia_path / "_docker" / "kube_config_inline" + self.helper._write_inline(inline_file) + self.assertTrue(inline_file.exists()) + mock_call.assert_called_once() - 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.helpers.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" / "kube_config_inline").exists() - ) + def test_create_inline_linux(self, mock_is_wsl, mock_call): + fake_data = {"clusters": [{"cluster": {"server": "127.0.0.1"}}]} + with patch.object( + self.helper, + "_write_inline", + side_effect=lambda f: self._write_kube_config(fake_data), + ): + self.helper.create_inline() + inline_file = self.gaia_path / "_docker" / "kube_config_inline" + self.assertTrue(inline_file.exists()) + data = yaml.safe_load(open(inline_file)) + self.assertEqual(data["clusters"][0]["cluster"]["server"], "127.0.0.1") + @patch("subprocess.call", return_value=0) @patch("gaiaflow.managers.helpers.is_wsl", return_value=False) - def test_create_inline_windows_branch(self, _): + def test_patch_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() + kube_config = self._write_kube_config({"clusters": [{"cluster": {"server": "127.0.0.1"}}]}) + helper._patch_kube_config(yaml.safe_load(open(kube_config)), kube_config) + data = yaml.safe_load(open(kube_config)) + cluster = data["clusters"][0]["cluster"] + self.assertEqual(cluster["server"], "host.docker.internal") + self.assertTrue(cluster["insecure-skip-tls-verify"]) - @patch("gaiaflow.managers.helpers.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("subprocess.call", return_value=0) @patch("gaiaflow.managers.helpers.is_wsl", return_value=True) - def test_patch_wsl(self, _): + def test_patch_wsl_branch(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"] + inline_file = self._write_kube_config( + {"clusters": [{"cluster": {"server": "localhost"}}]} + ) + helper._patch_kube_config(yaml.safe_load(open(inline_file)), inline_file) + data = yaml.safe_load(open(inline_file)) + cluster = data["clusters"][0]["cluster"] + self.assertEqual(cluster["server"], "https://192.168.49.2:8443") + self.assertTrue(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 TestBaseDockerHandler(unittest.TestCase): def setUp(self): From 2785223b12eecf337a14ecfce4adc05e9478e771 Mon Sep 17 00:00:00 2001 From: Yogesh Kumar Baljeet Singh Date: Tue, 16 Sep 2025 16:47:45 +0200 Subject: [PATCH 15/15] little cleanup --- src/gaiaflow/managers/utils.py | 7 ++++--- tests/managers/test_utils.py | 4 ++-- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/src/gaiaflow/managers/utils.py b/src/gaiaflow/managers/utils.py index dfa4480..5347379 100644 --- a/src/gaiaflow/managers/utils.py +++ b/src/gaiaflow/managers/utils.py @@ -175,7 +175,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() - log_info("state_file: " + str(state_file)) + log_info("found gaiaflow state file: " + str(state_file)) if not state_file.exists(): log_error( "State file not found at ~/.gaiaflow/state.json. Please run the services." @@ -186,13 +186,14 @@ def delete_project_state(gaiaflow_path: Path): with open(state_file, "r") as f: state = json.load(f) - log_info("found! " + str(state.get("gaiaflow_path")) + str(state)) + assert isinstance(state, dict) + 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, AttributeError, Exception): + except (json.JSONDecodeError, FileNotFoundError, AssertionError, Exception): raise diff --git a/tests/managers/test_utils.py b/tests/managers/test_utils.py index 7a586ad..2c4d72a 100644 --- a/tests/managers/test_utils.py +++ b/tests/managers/test_utils.py @@ -203,13 +203,13 @@ def test_state_file_missing(self): self.assertFalse(self.state_file.exists()) def test_delete_raises_jsondecodeerror(self): - self.state_file.write_text("{ invalid_data }") + self.state_file.write_text("{ invalid_data: '1' }") 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): + with self.assertRaises(AssertionError): utils.delete_project_state(self.gaiaflow_path) def test_update_project_state(self):