feat: support required_registration_fraction for secondary workers (#…

…11970) - [ ] Regenerate this pull request now. PiperOrigin-RevId: 578743167 Source-Link: googleapis/googleapis@a70dd17 Source-Link: googleapis/googleapis-gen@501cba3 Copy-Tag: eyJwIjoicGFja2FnZXMvZ29vZ2xlLWNsb3VkLWRhdGFwcm9jLy5Pd2xCb3QueWFtbCIsImgiOiI1MDFjYmEzMzM4NjI3YTgyMzU5ZDJmMzA2OWQ3ZTVhZjExYjdiM2RkIn0= --------- Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
googleapis · Nov 2, 2023 · 52d4558 · 52d4558
1 parent a27b282
commit 52d4558
Show file tree

Hide file tree

Showing 8 changed files with 53 additions and 0 deletions.
diff --git a/packages/google-cloud-dataproc/google/cloud/dataproc/__init__.py b/packages/google-cloud-dataproc/google/cloud/dataproc/__init__.py
@@ -114,6 +114,7 @@
     ShieldedInstanceConfig,
     SoftwareConfig,
     StartClusterRequest,
+    StartupConfig,
     StopClusterRequest,
     UpdateClusterRequest,
     VirtualClusterConfig,
@@ -268,6 +269,7 @@
     "ShieldedInstanceConfig",
     "SoftwareConfig",
     "StartClusterRequest",
+    "StartupConfig",
     "StopClusterRequest",
     "UpdateClusterRequest",
     "VirtualClusterConfig",

diff --git a/packages/google-cloud-dataproc/google/cloud/dataproc_v1/__init__.py b/packages/google-cloud-dataproc/google/cloud/dataproc_v1/__init__.py
@@ -98,6 +98,7 @@
     ShieldedInstanceConfig,
     SoftwareConfig,
     StartClusterRequest,
+    StartupConfig,
     StopClusterRequest,
     UpdateClusterRequest,
     VirtualClusterConfig,
@@ -305,6 +306,7 @@
     "SparkSqlBatch",
     "SparkSqlJob",
     "StartClusterRequest",
+    "StartupConfig",
     "StopClusterRequest",
     "SubmitJobRequest",
     "TemplateParameter",

diff --git a/packages/google-cloud-dataproc/google/cloud/dataproc_v1/types/__init__.py b/packages/google-cloud-dataproc/google/cloud/dataproc_v1/types/__init__.py
@@ -75,6 +75,7 @@
     ShieldedInstanceConfig,
     SoftwareConfig,
     StartClusterRequest,
+    StartupConfig,
     StopClusterRequest,
     UpdateClusterRequest,
     VirtualClusterConfig,
@@ -217,6 +218,7 @@
     "ShieldedInstanceConfig",
     "SoftwareConfig",
     "StartClusterRequest",
+    "StartupConfig",
     "StopClusterRequest",
     "UpdateClusterRequest",
     "VirtualClusterConfig",

diff --git a/packages/google-cloud-dataproc/google/cloud/dataproc_v1/types/clusters.py b/packages/google-cloud-dataproc/google/cloud/dataproc_v1/types/clusters.py
@@ -40,6 +40,7 @@
         "ShieldedInstanceConfig",
         "ConfidentialInstanceConfig",
         "InstanceGroupConfig",
+        "StartupConfig",
         "InstanceReference",
         "ManagedGroupConfig",
         "InstanceFlexibilityPolicy",
@@ -856,6 +857,10 @@ class InstanceGroupConfig(proto.Message):
             Optional. Instance flexibility Policy
             allowing a mixture of VM shapes and provisioning
             models.
+        startup_config (google.cloud.dataproc_v1.types.StartupConfig):
+            Optional. Configuration to handle the startup
+            of instances during cluster create and update
+            process.
     """
 
     class Preemptibility(proto.Enum):
@@ -952,6 +957,39 @@ class Preemptibility(proto.Enum):
         number=13,
         message="InstanceFlexibilityPolicy",
     )
+    startup_config: "StartupConfig" = proto.Field(
+        proto.MESSAGE,
+        number=14,
+        message="StartupConfig",
+    )
+
+
+class StartupConfig(proto.Message):
+    r"""Configuration to handle the startup of instances during
+    cluster create and update process.
+
+
+    .. _oneof: https://proto-plus-python.readthedocs.io/en/stable/fields.html#oneofs-mutually-exclusive-fields
+
+    Attributes:
+        required_registration_fraction (float):
+            Optional. The config setting to enable cluster creation/
+            updation to be successful only after
+            required_registration_fraction of instances are up and
+            running. This configuration is applicable to only secondary
+            workers for now. The cluster will fail if
+            required_registration_fraction of instances are not
+            available. This will include instance creation, agent
+            registration, and service registration (if enabled).
+
+            This field is a member of `oneof`_ ``_required_registration_fraction``.
+    """
+
+    required_registration_fraction: float = proto.Field(
+        proto.DOUBLE,
+        number=1,
+        optional=True,
+    )
 
 
 class InstanceReference(proto.Message):

diff --git a/packages/google-cloud-dataproc/noxfile.py b/packages/google-cloud-dataproc/noxfile.py
@@ -29,8 +29,10 @@
 
 BLACK_VERSION = "black[jupyter]==23.7.0"
 ISORT_VERSION = "isort==5.11.0"
+
 LINT_PATHS = ["docs", "google", "tests", "noxfile.py", "setup.py"]
 
+
 DEFAULT_PYTHON_VERSION = "3.9"
 
 UNIT_TEST_PYTHON_VERSIONS: List[str] = ["3.7", "3.8", "3.9", "3.10", "3.11"]
@@ -89,6 +91,7 @@ def lint(session):
         "--check",
         *LINT_PATHS,
     )
+
     session.run("flake8", "google", "tests")
 
 

diff --git a/packages/google-cloud-dataproc/tests/unit/gapic/dataproc_v1/test_cluster_controller.py b/packages/google-cloud-dataproc/tests/unit/gapic/dataproc_v1/test_cluster_controller.py
@@ -2868,6 +2868,7 @@ def test_create_cluster_rest(request_type):
                         {"machine_type": "machine_type_value", "vm_count": 875}
                     ],
                 },
+                "startup_config": {"required_registration_fraction": 0.3216},
             },
             "worker_config": {},
             "secondary_worker_config": {},
@@ -3437,6 +3438,7 @@ def test_update_cluster_rest(request_type):
                         {"machine_type": "machine_type_value", "vm_count": 875}
                     ],
                 },
+                "startup_config": {"required_registration_fraction": 0.3216},
             },
             "worker_config": {},
             "secondary_worker_config": {},

diff --git a/packages/google-cloud-dataproc/tests/unit/gapic/dataproc_v1/test_node_group_controller.py b/packages/google-cloud-dataproc/tests/unit/gapic/dataproc_v1/test_node_group_controller.py
@@ -1569,6 +1569,7 @@ def test_create_node_group_rest(request_type):
                     {"machine_type": "machine_type_value", "vm_count": 875}
                 ],
             },
+            "startup_config": {"required_registration_fraction": 0.3216},
         },
         "labels": {},
     }

diff --git a/...ages/google-cloud-dataproc/tests/unit/gapic/dataproc_v1/test_workflow_template_service.py b/...ages/google-cloud-dataproc/tests/unit/gapic/dataproc_v1/test_workflow_template_service.py
@@ -2834,6 +2834,7 @@ def test_create_workflow_template_rest(request_type):
                                 {"machine_type": "machine_type_value", "vm_count": 875}
                             ],
                         },
+                        "startup_config": {"required_registration_fraction": 0.3216},
                     },
                     "worker_config": {},
                     "secondary_worker_config": {},
@@ -4026,6 +4027,7 @@ def test_instantiate_inline_workflow_template_rest(request_type):
                                 {"machine_type": "machine_type_value", "vm_count": 875}
                             ],
                         },
+                        "startup_config": {"required_registration_fraction": 0.3216},
                     },
                     "worker_config": {},
                     "secondary_worker_config": {},
@@ -4658,6 +4660,7 @@ def test_update_workflow_template_rest(request_type):
                                 {"machine_type": "machine_type_value", "vm_count": 875}
                             ],
                         },
+                        "startup_config": {"required_registration_fraction": 0.3216},
                     },
                     "worker_config": {},
                     "secondary_worker_config": {},