diff --git a/CHANGELOG.md b/CHANGELOG.md index 84cf359..00a0fda 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,8 @@ ## [Unreleased] +- [Feature] Added distributed training support using CustomTrainingJobOp for scalable model training with worker replicas. Nodes can be configured for distributed training via node names or tags, supporting primary + worker pool architecture with configurable machine types and accelerators. + ## [0.12.0] - 2025-03-11 - Support for python 3.11 & 3.12 added, dropped support for python 3.8 diff --git a/kedro_vertexai/config.py b/kedro_vertexai/config.py index 93ed7c1..eacbee1 100644 --- a/kedro_vertexai/config.py +++ b/kedro_vertexai/config.py @@ -123,6 +123,38 @@ # allow_queueing: false # max_run_count: none # max_concurrent_run_count: 1 + + # Optional distributed training configuration + # distributed_training: + # # Enable distributed training for specific node names + # enabled_for_node_names: + # - "training_node" + # - "model_training" + # + # # Enable distributed training for nodes with specific tags + # enabled_for_tags: + # - "distributed" + # - "gpu-intensive" + # + # # Primary replica configuration (must have replica_count = 1) + # primary_pool: + # machine_type: "n1-standard-4" + # replica_count: 1 + # accelerator_type: "NVIDIA_TESLA_T4" + # accelerator_count: 1 + # + # # Worker pool configuration (can have replica_count > 1) + # worker_pool: + # machine_type: "n1-standard-4" + # replica_count: 2 + # accelerator_type: "NVIDIA_TESLA_T4" + # accelerator_count: 1 + # + # # Base output directory for distributed training jobs + # base_output_directory: "gs://your-bucket/distributed-training-output/" + # + # # Service account for distributed training (optional, defaults to global service_account) + # service_account: "distributed-training@my-project.iam.gserviceaccount.com" """ @@ -170,6 +202,8 @@ class GroupingConfig(BaseModel): def class_valid(cls, v, values, **kwargs): try: grouper_class = dynamic_load_class(v) + if grouper_class is None: + raise ValueError(f"Could not load grouping class {v}") class_sig = signature(grouper_class) if "params" in values.data: class_sig.bind(None, **values.data["params"]) @@ -225,6 +259,22 @@ class ScheduleConfig(BaseModel): max_concurrent_run_count: Optional[int] = 1 +class WorkerPoolConfig(BaseModel): + machine_type: str = "n1-standard-4" + replica_count: int = 1 + accelerator_type: Optional[str] = None + accelerator_count: Optional[int] = None + + +class DistributedTrainingConfig(BaseModel): + enabled_for_node_names: Optional[List[str]] = [] + enabled_for_tags: Optional[List[str]] = [] + primary_pool: Optional[WorkerPoolConfig] = WorkerPoolConfig() + worker_pool: Optional[WorkerPoolConfig] = WorkerPoolConfig(replica_count=2) + base_output_directory: Optional[str] = None + service_account: Optional[str] = None + + class RunConfig(BaseModel): image: str root: Optional[str] = None @@ -243,13 +293,35 @@ class RunConfig(BaseModel): dynamic_config_providers: Optional[List[DynamicConfigProviderConfig]] = [] mlflow: Optional[MLFlowVertexAIConfig] = None schedules: Optional[Dict[str, ScheduleConfig]] = None + distributed_training: Optional[DistributedTrainingConfig] = None def resources_for(self, node: str, tags: Optional[set] = None): + if self.resources is None: + return {} default_config = self.resources["__default__"].dict() - return self._config_for(node, tags, self.resources, default_config) + return self._config_for(node, tags or set(), self.resources, default_config) def node_selectors_for(self, node: str, tags: Optional[set] = None): - return self._config_for(node, tags, self.node_selectors) + if self.node_selectors is None: + return {} + return self._config_for(node, tags or set(), self.node_selectors) + + def should_use_distributed_training(self, node: str, tags: Optional[set] = None) -> bool: + """Check if a node should use distributed training based on configuration.""" + if not self.distributed_training: + return False + + tags = tags or set() + + # Check node names + if node in (self.distributed_training.enabled_for_node_names or []): + return True + + # Check tags + if any(tag in (self.distributed_training.enabled_for_tags or []) for tag in tags): + return True + + return False @staticmethod def _config_for( diff --git a/kedro_vertexai/generator.py b/kedro_vertexai/generator.py index 6e91217..61dd7a8 100644 --- a/kedro_vertexai/generator.py +++ b/kedro_vertexai/generator.py @@ -10,6 +10,7 @@ from kfp import dsl from kfp.dsl import PipelineTask from makefun import with_signature +from google_cloud_pipeline_components.v1.custom_job import CustomTrainingJobOp from kedro_vertexai.config import ( KedroVertexAIRunnerConfig, @@ -201,25 +202,41 @@ def _build_kfp_tasks( ] ).strip() - @dsl.container_component - @with_signature(f"{name.replace('-', '_')}({params_signature})") - def component(*args, **kwargs): - dynamic_parameters = ",".join( - [f"{k}={kwargs[k]}" for k in params.keys()] - ) - - return dsl.ContainerSpec( + # Check if this node should use distributed training + if self.run_config.should_use_distributed_training(group_name, tags): + # Create CustomTrainingJobOp for distributed training + task = self._create_custom_training_job_task( + name=name, image=image, - command=["/bin/bash", "-c"], - args=[ - node_command, - " --params", # TODO what if there is no dynamic params? - f" {dynamic_parameters}", - ], + kedro_command=kedro_command, + nodes_group=nodes_group, + tags=tags, + params_signature=params_signature, + component_params=component_params, + should_add_params=should_add_params, ) - - task = component(**component_params) - self._configure_resources(name, tags, task) + else: + # Create standard container component + @dsl.container_component + @with_signature(f"{name.replace('-', '_')}({params_signature})") + def component(*args, **kwargs): + dynamic_parameters = ",".join( + [f"{k}={kwargs[k]}" for k in params.keys()] + ) + + return dsl.ContainerSpec( + image=image, + command=["/bin/bash", "-c"], + args=[ + node_command, + " --params", # TODO what if there is no dynamic params? + f" {dynamic_parameters}", + ], + ) + + task = component(**component_params) + self._configure_resources(name, tags, task) + kfp_tasks[name] = task return kfp_tasks @@ -249,6 +266,114 @@ def _generate_gcp_env_vars_command(self) -> str: region = vertex_conf.get("region") return f"GCP_PROJECT_ID={project_id} GCP_REGION={region}" + def _create_custom_training_job_task( + self, + name: str, + image: str, + kedro_command: str, + nodes_group: List, + tags: set, + params_signature: str, + component_params: Dict, + should_add_params: bool, + ) -> PipelineTask: + """Create a CustomTrainingJobOp task for distributed training.""" + + if not self.run_config.distributed_training: + raise ValueError("Distributed training config is required for CustomTrainingJobOp") + + dt_config = self.run_config.distributed_training + + # Ensure primary_pool and worker_pool are not None + if not dt_config.primary_pool: + raise ValueError("Primary pool configuration is required for distributed training") + if not dt_config.worker_pool: + raise ValueError("Worker pool configuration is required for distributed training") + + # Build the full command with all necessary setup + full_command = " ".join([ + h + " " if (h := self._generate_hosts_file()) else "", + self._generate_params_command(should_add_params), + "MLFLOW_RUN_ID=\"{{$.inputs.parameters['mlflow_run_id']}}\" " + if is_mlflow_enabled() + else "", + self._generate_gcp_env_vars_command(), + kedro_command, + ]).strip() + + # Build worker pool specs based on configuration + worker_pool_specs = [] + + # Get resource configuration from existing resources config + resources = self.run_config.resources_for(name, tags) + + # Build primary spec with machine type from config or default + primary_machine_type = dt_config.primary_pool.machine_type + primary_spec = { + "machine_spec": { + "machine_type": primary_machine_type, + }, + "replica_count": 1, # Primary must always be 1 + "container_spec": { + "image_uri": image, + "command": ["/bin/bash", "-c"], + "args": [full_command], + }, + } + + # Add accelerator config from distributed training config + if dt_config.primary_pool.accelerator_type: + primary_spec["machine_spec"]["accelerator_type"] = dt_config.primary_pool.accelerator_type + if dt_config.primary_pool.accelerator_count: + primary_spec["machine_spec"]["accelerator_count"] = dt_config.primary_pool.accelerator_count + + worker_pool_specs.append(primary_spec) + + # Worker pool (can have replica_count > 1) + if dt_config.worker_pool.replica_count > 0: + worker_machine_type = dt_config.worker_pool.machine_type + worker_spec = { + "machine_spec": { + "machine_type": worker_machine_type, + }, + "replica_count": dt_config.worker_pool.replica_count, + "container_spec": { + "image_uri": image, + "command": ["/bin/bash", "-c"], + "args": [full_command], + }, + } + + # Add accelerator config from distributed training config + if dt_config.worker_pool.accelerator_type: + worker_spec["machine_spec"]["accelerator_type"] = dt_config.worker_pool.accelerator_type + if dt_config.worker_pool.accelerator_count: + worker_spec["machine_spec"]["accelerator_count"] = dt_config.worker_pool.accelerator_count + + worker_pool_specs.append(worker_spec) + + # Create CustomTrainingJobOp task + # Note: While KFP compiler generates names like "comp-custom-training-job", + # "comp-custom-training-job-2", etc., creating unique function instances + # ensures each distributed training node gets processed separately. + + # Use service account from distributed training config, fallback to global service account + service_account = dt_config.service_account or self.run_config.service_account or "" + + # Create the task directly using CustomTrainingJobOp + task = CustomTrainingJobOp( + display_name=f"distributed-{name}", + worker_pool_specs=worker_pool_specs, + base_output_directory=dt_config.base_output_directory or f"gs://{self.run_config.root}/distributed-training-output/", + service_account=service_account, + **component_params + ) + + # Set display name to the node name for better identification on UI + task.set_display_name(name) + + return task + def _configure_resources(self, name: str, tags: set, task: PipelineTask): resources = self.run_config.resources_for(name, tags) node_selectors = self.run_config.node_selectors_for(name, tags) diff --git a/poetry.lock b/poetry.lock index 46e2d3c..97dd753 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 2.1.1 and should not be changed by hand. +# This file is automatically @generated by Poetry 2.1.3 and should not be changed by hand. [[package]] name = "aiohappyeyeballs" @@ -1272,12 +1272,12 @@ files = [ google-auth = ">=2.14.1,<3.0.0" googleapis-common-protos = ">=1.56.2,<2.0.0" grpcio = [ - {version = ">=1.33.2,<2.0dev", optional = true, markers = "extra == \"grpc\""}, {version = ">=1.49.1,<2.0dev", optional = true, markers = "python_version >= \"3.11\" and extra == \"grpc\""}, + {version = ">=1.33.2,<2.0dev", optional = true, markers = "python_version < \"3.11\" and extra == \"grpc\""}, ] grpcio-status = [ - {version = ">=1.33.2,<2.0.dev0", optional = true, markers = "extra == \"grpc\""}, {version = ">=1.49.1,<2.0.dev0", optional = true, markers = "python_version >= \"3.11\" and extra == \"grpc\""}, + {version = ">=1.33.2,<2.0.dev0", optional = true, markers = "extra == \"grpc\""}, ] proto-plus = ">=1.22.3,<2.0.0" protobuf = ">=3.19.5,<3.20.0 || >3.20.0,<3.20.1 || >3.20.1,<4.21.0 || >4.21.0,<4.21.1 || >4.21.1,<4.21.2 || >4.21.2,<4.21.3 || >4.21.3,<4.21.4 || >4.21.4,<4.21.5 || >4.21.5,<7.0.0" @@ -1463,6 +1463,27 @@ grpc-google-iam-v1 = ">=0.12.4,<1.0.0" proto-plus = ">=1.22.3,<2.0.0" protobuf = ">=3.20.2,<4.21.0 || >4.21.0,<4.21.1 || >4.21.1,<4.21.2 || >4.21.2,<4.21.3 || >4.21.3,<4.21.4 || >4.21.4,<4.21.5 || >4.21.5,<7.0.0" +[[package]] +name = "google-cloud-pipeline-components" +version = "2.20.1" +description = "This SDK enables a set of First Party (Google owned) pipeline components that allow users to take their experience from Vertex AI SDK and other Google Cloud services and create a corresponding pipeline using KFP or Managed Pipelines." +optional = false +python-versions = ">=3.8.0" +groups = ["main"] +files = [ + {file = "google_cloud_pipeline_components-2.20.1-py3-none-any.whl", hash = "sha256:77b3b735dc13f6868e3ed9281af3f195a1becbae0960755bb142286d784a06bd"}, +] + +[package.dependencies] +google-api-core = ">=1.31.5,<2.0.dev0 || >2.3.0,<3.0.0dev" +google-cloud-aiplatform = ">=1.14.0,<2" +Jinja2 = ">=3.1.2,<4" +kfp = ">=2.6.0,<3.0.0" + +[package.extras] +docs = ["autodocsumm (==0.2.9)", "commonmark (==0.9.1)", "grpcio-status (<=1.47.0)", "m2r2 (==0.3.3.post2)", "protobuf (>=4.21.1,<5)", "sphinx (>=5.0.2,<6.0.0)", "sphinx-immaterial (==0.9.0)", "sphinx-notfound-page (==0.8.3)", "sphinx-rtd-theme (==2.0.0)"] +tests = ["flake8 (>=3.0.0)", "mock (>=4.0.0)", "pytest (>=6.0.0)"] + [[package]] name = "google-cloud-resource-manager" version = "1.14.1" @@ -3531,9 +3552,9 @@ files = [ [package.dependencies] numpy = [ - {version = ">=1.22.4", markers = "python_version < \"3.11\""}, {version = ">=1.23.2", markers = "python_version == \"3.11\""}, {version = ">=1.26.0", markers = "python_version >= \"3.12\""}, + {version = ">=1.22.4", markers = "python_version < \"3.11\""}, ] python-dateutil = ">=2.8.2" pytz = ">=2020.1" @@ -5591,4 +5612,4 @@ mlflow = ["kedro-mlflow"] [metadata] lock-version = "2.1" python-versions = ">=3.9,<3.13" -content-hash = "6e965ede43573a8ab9c89452a860f8c526b5d09ca92d103c78e5636594f54c22" +content-hash = "c9905be2e190d7eb23f9d59b87cae8ee116dde12414018cb43483053ed5c93f6" diff --git a/pyproject.toml b/pyproject.toml index e098569..b0d1736 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,6 +44,7 @@ pydantic = ">=2,<3" google-auth = "<3" google-cloud-scheduler = ">=2.3.2" google-cloud-iam = "<3" +google-cloud-pipeline-components = ">=2.20.1" gcsfs = ">=2022.1" fsspec = ">=2022.1" google-cloud-storage = "<3.0.0" diff --git a/tests/test_distributed_training.py b/tests/test_distributed_training.py new file mode 100644 index 0000000..6505d71 --- /dev/null +++ b/tests/test_distributed_training.py @@ -0,0 +1,649 @@ +"""Test distributed training configuration and generation""" + +import unittest +from tempfile import NamedTemporaryFile +from unittest.mock import MagicMock, patch + +import kfp +import yaml +from kedro.pipeline import Pipeline, node +from kfp.compiler import Compiler +from pydantic import ValidationError + +from kedro_vertexai.config import ( + DistributedTrainingConfig, + PluginConfig, + WorkerPoolConfig, +) +from kedro_vertexai.generator import PipelineGenerator +from tests.utils import environment + + +def dummy_train_func(input_data: str) -> str: + return input_data # pragma: no cover + + +def dummy_preprocess_func(input_data: str) -> str: + return input_data # pragma: no cover + + +class TestDistributedTrainingConfig(unittest.TestCase): + """Test distributed training configuration classes""" + + def test_worker_pool_config_defaults(self): + """Test WorkerPoolConfig with default values""" + config = WorkerPoolConfig() + self.assertEqual(config.machine_type, "n1-standard-4") + self.assertEqual(config.replica_count, 1) + self.assertIsNone(config.accelerator_type) + self.assertIsNone(config.accelerator_count) + + def test_worker_pool_config_custom_values(self): + """Test WorkerPoolConfig with custom values""" + config = WorkerPoolConfig( + machine_type="n1-standard-8", + replica_count=4, + accelerator_type="NVIDIA_TESLA_T4", + accelerator_count=2, + ) + self.assertEqual(config.machine_type, "n1-standard-8") + self.assertEqual(config.replica_count, 4) + self.assertEqual(config.accelerator_type, "NVIDIA_TESLA_T4") + self.assertEqual(config.accelerator_count, 2) + + def test_distributed_training_config_defaults(self): + """Test DistributedTrainingConfig with default values""" + config = DistributedTrainingConfig() + self.assertEqual(config.enabled_for_node_names, []) + self.assertEqual(config.enabled_for_tags, []) + + # Test primary pool defaults + primary_pool = config.primary_pool + self.assertIsNotNone(primary_pool) + if primary_pool: + self.assertEqual(primary_pool.replica_count, 1) + + # Test worker pool defaults + worker_pool = config.worker_pool + self.assertIsNotNone(worker_pool) + if worker_pool: + self.assertEqual(worker_pool.replica_count, 2) + + self.assertIsNone(config.base_output_directory) + self.assertIsNone(config.service_account) + + def test_distributed_training_config_custom_values(self): + """Test DistributedTrainingConfig with custom values""" + config = DistributedTrainingConfig( + enabled_for_node_names=["train_model", "train_embedding"], + enabled_for_tags=["distributed", "gpu"], + primary_pool=WorkerPoolConfig( + machine_type="n1-standard-8", accelerator_type="NVIDIA_TESLA_T4" + ), + worker_pool=WorkerPoolConfig( + machine_type="n1-standard-8", replica_count=4 + ), + base_output_directory="gs://my-bucket/output/", + service_account="distributed-training@my-project.iam.gserviceaccount.com", + ) + self.assertEqual(config.enabled_for_node_names, ["train_model", "train_embedding"]) + self.assertEqual(config.enabled_for_tags, ["distributed", "gpu"]) + + # Test primary pool configuration + primary_pool = config.primary_pool + if primary_pool: + self.assertEqual(primary_pool.machine_type, "n1-standard-8") + + # Test worker pool configuration + worker_pool = config.worker_pool + if worker_pool: + self.assertEqual(worker_pool.replica_count, 4) + + self.assertEqual(config.base_output_directory, "gs://my-bucket/output/") + self.assertEqual(config.service_account, "distributed-training@my-project.iam.gserviceaccount.com") + + def test_should_use_distributed_training_node_names(self): + """Test should_use_distributed_training with node names""" + config_yaml = """ +project_id: test-project +region: test-region +run_config: + image: test-image + experiment_name: test-experiment + distributed_training: + enabled_for_node_names: + - train_model + - train_embedding +""" + config = PluginConfig.model_validate(yaml.safe_load(config_yaml)) + + # Test node names that should use distributed training + self.assertTrue(config.run_config.should_use_distributed_training("train_model")) + self.assertTrue(config.run_config.should_use_distributed_training("train_embedding")) + + # Test node names that should not use distributed training + self.assertFalse(config.run_config.should_use_distributed_training("preprocess")) + self.assertFalse(config.run_config.should_use_distributed_training("evaluate")) + + def test_should_use_distributed_training_tags(self): + """Test should_use_distributed_training with tags""" + config_yaml = """ +project_id: test-project +region: test-region +run_config: + image: test-image + experiment_name: test-experiment + distributed_training: + enabled_for_tags: + - distributed + - gpu-intensive +""" + config = PluginConfig.model_validate(yaml.safe_load(config_yaml)) + + # Test tags that should use distributed training + self.assertTrue(config.run_config.should_use_distributed_training("any_node", {"distributed"})) + self.assertTrue(config.run_config.should_use_distributed_training("any_node", {"gpu-intensive"})) + self.assertTrue(config.run_config.should_use_distributed_training("any_node", {"distributed", "other"})) + + # Test tags that should not use distributed training + self.assertFalse(config.run_config.should_use_distributed_training("any_node", {"standard"})) + self.assertFalse(config.run_config.should_use_distributed_training("any_node", {"cpu-only"})) + self.assertFalse(config.run_config.should_use_distributed_training("any_node", set())) + + def test_should_use_distributed_training_mixed(self): + """Test should_use_distributed_training with both node names and tags""" + config_yaml = """ +project_id: test-project +region: test-region +run_config: + image: test-image + experiment_name: test-experiment + distributed_training: + enabled_for_node_names: + - train_model + enabled_for_tags: + - distributed +""" + config = PluginConfig.model_validate(yaml.safe_load(config_yaml)) + + # Test node name match + self.assertTrue(config.run_config.should_use_distributed_training("train_model")) + + # Test tag match + self.assertTrue(config.run_config.should_use_distributed_training("other_node", {"distributed"})) + + # Test no match + self.assertFalse(config.run_config.should_use_distributed_training("other_node", {"standard"})) + + def test_should_use_distributed_training_disabled(self): + """Test should_use_distributed_training when distributed training is not configured""" + config_yaml = """ +project_id: test-project +region: test-region +run_config: + image: test-image + experiment_name: test-experiment +""" + config = PluginConfig.model_validate(yaml.safe_load(config_yaml)) + + # Should return False when distributed training is not configured + self.assertFalse(config.run_config.should_use_distributed_training("any_node")) + self.assertFalse(config.run_config.should_use_distributed_training("any_node", {"any_tag"})) + + def test_plugin_config_with_distributed_training(self): + """Test PluginConfig with distributed training configuration""" + config_yaml = """ +project_id: test-project +region: test-region +run_config: + image: test-image + experiment_name: test-experiment + distributed_training: + enabled_for_node_names: + - train_model + enabled_for_tags: + - distributed + primary_pool: + machine_type: n1-standard-8 + replica_count: 1 + accelerator_type: NVIDIA_TESLA_T4 + accelerator_count: 1 + worker_pool: + machine_type: n1-standard-8 + replica_count: 4 + accelerator_type: NVIDIA_TESLA_T4 + accelerator_count: 1 + base_output_directory: gs://my-bucket/output/ +""" + config = PluginConfig.model_validate(yaml.safe_load(config_yaml)) + + dt_config = config.run_config.distributed_training + self.assertIsNotNone(dt_config) + + if dt_config: + self.assertEqual(dt_config.enabled_for_node_names, ["train_model"]) + self.assertEqual(dt_config.enabled_for_tags, ["distributed"]) + + primary_pool = dt_config.primary_pool + if primary_pool: + self.assertEqual(primary_pool.machine_type, "n1-standard-8") + self.assertEqual(primary_pool.accelerator_type, "NVIDIA_TESLA_T4") + + worker_pool = dt_config.worker_pool + if worker_pool: + self.assertEqual(worker_pool.replica_count, 4) + + self.assertEqual(dt_config.base_output_directory, "gs://my-bucket/output/") + + +class TestDistributedTrainingGenerator(unittest.TestCase): + """Test distributed training pipeline generation""" + + def create_pipeline(self): + """Create a test pipeline with distributed and standard nodes""" + return Pipeline( + [ + node( + dummy_preprocess_func, + "raw_data", + "preprocessed_data", + name="preprocess", + tags=["preprocessing"], + ), + node( + dummy_train_func, + "preprocessed_data", + "model", + name="train_model", + tags=["training", "distributed"], + ), + node( + dummy_train_func, + "preprocessed_data", + "embeddings", + name="train_embedding", + tags=["training"], + ), + ] + ) + + def create_generator(self, config={}, params={}, catalog={}): + """Create a PipelineGenerator for testing""" + project_name = "test-distributed-training" + config_loader = MagicMock() + config_loader.get.return_value = catalog + context = type( + "obj", + (object,), + { + "env": "unittests", + "params": params, + "config_loader": config_loader, + }, + ) + + self.pipelines_under_test = {"pipeline": self.create_pipeline()} + + config_with_defaults = { + "image": "test-image", + "root": "test-bucket/test-suffix", + "experiment_name": "test-experiment", + "run_name": "test-run", + } + config_with_defaults.update(config) + + self.generator_under_test = PipelineGenerator( + PluginConfig.model_validate( + { + "project_id": "test-project", + "region": "test-region", + "run_config": config_with_defaults, + } + ), + project_name, + context, + "test-run-name", + ) + + def test_should_generate_standard_container_when_distributed_training_disabled(self): + """Test that standard container components are generated when distributed training is disabled""" + self.create_generator() + + with patch("kedro.framework.project.pipelines", new=self.pipelines_under_test): + pipeline = self.generator_under_test.generate_pipeline( + "pipeline", "test-image", "test-token" + ) + + with NamedTemporaryFile(mode="rt", prefix="pipeline", suffix=".yaml") as spec_output: + Compiler().compile(pipeline, spec_output.name) + with open(spec_output.name) as f: + pipeline_spec = yaml.safe_load(f) + + # All nodes should be standard container components + executors = pipeline_spec["deploymentSpec"]["executors"] + self.assertIn("exec-preprocess", executors) + self.assertIn("exec-train-model", executors) + self.assertIn("exec-train-embedding", executors) + + # Check that they are container components (not CustomTrainingJobOp) + for executor_name in ["exec-preprocess", "exec-train-model", "exec-train-embedding"]: + self.assertIn("container", executors[executor_name]) + self.assertIn("args", executors[executor_name]["container"]) + + @patch("kedro_vertexai.generator.CustomTrainingJobOp") + def test_should_generate_custom_training_job_when_enabled_by_node_name(self, mock_custom_training_job): + """Test that CustomTrainingJobOp is generated for nodes enabled by name""" + config = { + "distributed_training": { + "enabled_for_node_names": ["train_model"], + "primary_pool": { + "machine_type": "n1-standard-8", + "replica_count": 1, + "accelerator_type": "NVIDIA_TESLA_T4", + "accelerator_count": 1, + }, + "worker_pool": { + "machine_type": "n1-standard-8", + "replica_count": 2, + "accelerator_type": "NVIDIA_TESLA_T4", + "accelerator_count": 1, + }, + "base_output_directory": "gs://test-bucket/output/", + } + } + + self.create_generator(config=config) + + with patch("kedro.framework.project.pipelines", new=self.pipelines_under_test): + pipeline = self.generator_under_test.generate_pipeline( + "pipeline", "test-image", "test-token" + ) + + # CustomTrainingJobOp should be called for the distributed training node + mock_custom_training_job.assert_called() + + # Check the call arguments + call_args = mock_custom_training_job.call_args + self.assertEqual(call_args[1]["display_name"], "distributed-train-model") + self.assertEqual(call_args[1]["base_output_directory"], "gs://test-bucket/output/") + + # Check worker pool specs + worker_pool_specs = call_args[1]["worker_pool_specs"] + self.assertEqual(len(worker_pool_specs), 2) # Primary + worker pool + + # Check primary pool + primary_pool = worker_pool_specs[0] + self.assertEqual(primary_pool["replica_count"], 1) + self.assertEqual(primary_pool["machine_spec"]["machine_type"], "n1-standard-8") + self.assertEqual(primary_pool["machine_spec"]["accelerator_type"], "NVIDIA_TESLA_T4") + self.assertEqual(primary_pool["machine_spec"]["accelerator_count"], 1) + + # Check worker pool + worker_pool = worker_pool_specs[1] + self.assertEqual(worker_pool["replica_count"], 2) + self.assertEqual(worker_pool["machine_spec"]["machine_type"], "n1-standard-8") + self.assertEqual(worker_pool["machine_spec"]["accelerator_type"], "NVIDIA_TESLA_T4") + self.assertEqual(worker_pool["machine_spec"]["accelerator_count"], 1) + + @patch("kedro_vertexai.generator.CustomTrainingJobOp") + def test_should_generate_custom_training_job_when_enabled_by_tag(self, mock_custom_training_job): + """Test that CustomTrainingJobOp is generated for nodes enabled by tag""" + config = { + "distributed_training": { + "enabled_for_tags": ["distributed"], + "primary_pool": { + "machine_type": "n1-standard-4", + "replica_count": 1, + }, + "worker_pool": { + "machine_type": "n1-standard-4", + "replica_count": 3, + }, + "base_output_directory": "gs://test-bucket/output/", + } + } + + self.create_generator(config=config) + + with patch("kedro.framework.project.pipelines", new=self.pipelines_under_test): + pipeline = self.generator_under_test.generate_pipeline( + "pipeline", "test-image", "test-token" + ) + + # CustomTrainingJobOp should be called for the node with "distributed" tag + mock_custom_training_job.assert_called() + + # Check that the correct node is using distributed training + call_args = mock_custom_training_job.call_args + self.assertEqual(call_args[1]["display_name"], "distributed-train-model") + + @patch("kedro_vertexai.generator.CustomTrainingJobOp") + def test_should_generate_mixed_components_when_partially_enabled(self, mock_custom_training_job): + """Test that both standard and distributed components are generated when partially enabled""" + config = { + "distributed_training": { + "enabled_for_node_names": ["train_model"], + "primary_pool": { + "machine_type": "n1-standard-4", + "replica_count": 1, + }, + "worker_pool": { + "machine_type": "n1-standard-4", + "replica_count": 2, + }, + "base_output_directory": "gs://test-bucket/output/", + } + } + + self.create_generator(config=config) + + with patch("kedro.framework.project.pipelines", new=self.pipelines_under_test): + pipeline = self.generator_under_test.generate_pipeline( + "pipeline", "test-image", "test-token" + ) + + with NamedTemporaryFile(mode="rt", prefix="pipeline", suffix=".yaml") as spec_output: + Compiler().compile(pipeline, spec_output.name) + with open(spec_output.name) as f: + pipeline_spec = yaml.safe_load(f) + + # Check that some nodes are standard containers + executors = pipeline_spec["deploymentSpec"]["executors"] + self.assertIn("exec-preprocess", executors) + self.assertIn("exec-train-embedding", executors) + + # These should be standard container components + self.assertIn("container", executors["exec-preprocess"]) + self.assertIn("container", executors["exec-train-embedding"]) + + # CustomTrainingJobOp should be called once for train_model + mock_custom_training_job.assert_called_once() + + def test_should_use_default_output_directory_when_not_specified(self): + """Test that default output directory is used when not specified""" + config = { + "distributed_training": { + "enabled_for_node_names": ["train_model"], + } + } + + self.create_generator(config=config) + + with patch("kedro_vertexai.generator.CustomTrainingJobOp") as mock_custom_training_job: + with patch("kedro.framework.project.pipelines", new=self.pipelines_under_test): + pipeline = self.generator_under_test.generate_pipeline( + "pipeline", "test-image", "test-token" + ) + + # Check that default output directory is used + call_args = mock_custom_training_job.call_args + expected_output_dir = "gs://test-bucket/test-suffix/distributed-training-output/" + self.assertEqual(call_args[1]["base_output_directory"], expected_output_dir) + + def test_should_handle_worker_pool_with_zero_replicas(self): + """Test that worker pool is omitted when replica count is 0""" + config = { + "distributed_training": { + "enabled_for_node_names": ["train_model"], + "primary_pool": { + "machine_type": "n1-standard-4", + "replica_count": 1, + }, + "worker_pool": { + "machine_type": "n1-standard-4", + "replica_count": 0, # No worker replicas + }, + } + } + + self.create_generator(config=config) + + with patch("kedro_vertexai.generator.CustomTrainingJobOp") as mock_custom_training_job: + with patch("kedro.framework.project.pipelines", new=self.pipelines_under_test): + pipeline = self.generator_under_test.generate_pipeline( + "pipeline", "test-image", "test-token" + ) + + # Check that only primary pool is included + call_args = mock_custom_training_job.call_args + worker_pool_specs = call_args[1]["worker_pool_specs"] + self.assertEqual(len(worker_pool_specs), 1) # Only primary pool + self.assertEqual(worker_pool_specs[0]["replica_count"], 1) + + @patch("kedro_vertexai.generator.CustomTrainingJobOp") + def test_should_use_service_account_from_distributed_training_config(self, mock_custom_training_job): + """Test that service account is passed from distributed training config""" + config = { + "service_account": "global@my-project.iam.gserviceaccount.com", + "distributed_training": { + "enabled_for_node_names": ["train_model"], + "service_account": "distributed@my-project.iam.gserviceaccount.com", + } + } + + self.create_generator(config=config) + + with patch("kedro.framework.project.pipelines", new=self.pipelines_under_test): + pipeline = self.generator_under_test.generate_pipeline( + "pipeline", "test-image", "test-token" + ) + + # Verify CustomTrainingJobOp was called with distributed training service account + mock_custom_training_job.assert_called_once() + call_args = mock_custom_training_job.call_args + self.assertEqual(call_args[1]["service_account"], "distributed@my-project.iam.gserviceaccount.com") + + @patch("kedro_vertexai.generator.CustomTrainingJobOp") + def test_should_fallback_to_global_service_account_when_not_specified(self, mock_custom_training_job): + """Test that service account falls back to global when not specified in distributed training config""" + config = { + "service_account": "global@my-project.iam.gserviceaccount.com", + "distributed_training": { + "enabled_for_node_names": ["train_model"], + } + } + + self.create_generator(config=config) + + with patch("kedro.framework.project.pipelines", new=self.pipelines_under_test): + pipeline = self.generator_under_test.generate_pipeline( + "pipeline", "test-image", "test-token" + ) + + # Verify CustomTrainingJobOp was called with global service account + mock_custom_training_job.assert_called_once() + call_args = mock_custom_training_job.call_args + self.assertEqual(call_args[1]["service_account"], "global@my-project.iam.gserviceaccount.com") + + @patch("kedro_vertexai.generator.CustomTrainingJobOp") + def test_should_use_empty_string_when_no_service_account_specified(self, mock_custom_training_job): + """Test that empty string is used when no service account is specified""" + config = { + "distributed_training": { + "enabled_for_node_names": ["train_model"], + } + } + + self.create_generator(config=config) + + with patch("kedro.framework.project.pipelines", new=self.pipelines_under_test): + pipeline = self.generator_under_test.generate_pipeline( + "pipeline", "test-image", "test-token" + ) + + # Verify CustomTrainingJobOp was called with empty string + mock_custom_training_job.assert_called_once() + call_args = mock_custom_training_job.call_args + self.assertEqual(call_args[1]["service_account"], "") + + def test_should_generate_unique_component_names_for_multiple_distributed_nodes(self): + """Test that multiple distributed training nodes get unique component names""" + config = { + "distributed_training": { + "enabled_for_node_names": ["train_model", "train_embedding"], + "primary_pool": { + "machine_type": "n1-standard-4", + "replica_count": 1, + }, + "worker_pool": { + "machine_type": "n1-standard-4", + "replica_count": 2, + }, + "base_output_directory": "gs://test-bucket/output/", + } + } + + self.create_generator(config=config) + + with patch("kedro.framework.project.pipelines", new=self.pipelines_under_test): + pipeline = self.generator_under_test.generate_pipeline( + "pipeline", "test-image", "test-token" + ) + + with NamedTemporaryFile(mode="rt", prefix="pipeline", suffix=".yaml") as spec_output: + Compiler().compile(pipeline, spec_output.name) + with open(spec_output.name) as f: + pipeline_spec = yaml.safe_load(f) + + # Check that both distributed training nodes have unique component names + components = pipeline_spec["components"] + component_names = list(components.keys()) + + # Check that the executors also have unique names + executors = pipeline_spec["deploymentSpec"]["executors"] + executor_names = list(executors.keys()) + + # Should have multiple custom training job components (one for each distributed node) + custom_training_components = [name for name in component_names if "custom-training-job" in name] + custom_training_executors = [name for name in executor_names if "custom-training-job" in name] + + # Should have exactly 2 custom training job components (train_model and train_embedding) + self.assertEqual(len(custom_training_components), 2, + f"Expected 2 custom training components, got {len(custom_training_components)}: {custom_training_components}") + + # Should have exactly 2 custom training job executors + self.assertEqual(len(custom_training_executors), 2, + f"Expected 2 custom training executors, got {len(custom_training_executors)}: {custom_training_executors}") + + # Component names should be unique + self.assertEqual(len(set(custom_training_components)), len(custom_training_components), + f"Component names should be unique: {custom_training_components}") + + # Executor names should be unique + self.assertEqual(len(set(custom_training_executors)), len(custom_training_executors), + f"Executor names should be unique: {custom_training_executors}") + + # Should also have one standard container component (preprocess) + standard_components = [name for name in component_names if "preprocess" in name] + self.assertEqual(len(standard_components), 1, + f"Expected 1 standard component for preprocess, got {len(standard_components)}: {standard_components}") + + # Verify the names are actually different + self.assertNotEqual(custom_training_components[0], custom_training_components[1], + "Custom training components should have different names") + self.assertNotEqual(custom_training_executors[0], custom_training_executors[1], + "Custom training executors should have different names") + + +if __name__ == "__main__": + unittest.main() \ No newline at end of file