Fix lint and nox

Author: ji-yaqi

google/cloud/aiplatform/hyperparameter_tuning.py

@@ -101,10 +101,7 @@ class DoubleParameterSpec(_ParameterSpec):
     _parameter_spec_value_key = "double_value_spec"
 
     def __init__(
-        self,
-        min: float,
-        max: float,
-        scale: str,
+        self, min: float, max: float, scale: str,
     ):
         """
         Value specification for a parameter in ``DOUBLE`` type.
@@ -138,10 +135,7 @@ class IntegerParameterSpec(_ParameterSpec):
     _parameter_spec_value_key = "integer_value_spec"
 
     def __init__(
-        self,
-        min: int,
-        max: int,
-        scale: str,
+        self, min: int, max: int, scale: str,
     ):
         """
         Value specification for a parameter in ``INTEGER`` type.
@@ -175,8 +169,7 @@ class CategoricalParameterSpec(_ParameterSpec):
     _parameter_spec_value_key = "categorical_value_spec"
 
     def __init__(
-        self,
-        values: Sequence[str],
+        self, values: Sequence[str],
     ):
         """Value specification for a parameter in ``CATEGORICAL`` type.
 
@@ -199,9 +192,7 @@ class DiscreteParameterSpec(_ParameterSpec):
     _parameter_spec_value_key = "discrete_value_spec"
 
     def __init__(
-        self,
-        values: Sequence[float],
-        scale: str,
+        self, values: Sequence[float], scale: str,
     ):
         """Value specification for a parameter in ``DISCRETE`` type.
 

google/cloud/aiplatform/jobs.py

@@ -19,7 +19,6 @@
 
 import abc
 import copy
-import sys
 import time
 
 from google.cloud import storage
@@ -568,11 +567,9 @@ def create(
         gapic_batch_prediction_job.output_config = output_config
 
         # Optional Fields
-        gapic_batch_prediction_job.encryption_spec = (
-            initializer.global_config.get_encryption_spec(
-                encryption_spec_key_name=encryption_spec_key_name,
-                select_version=select_version,
-            )
+        gapic_batch_prediction_job.encryption_spec = initializer.global_config.get_encryption_spec(
+            encryption_spec_key_name=encryption_spec_key_name,
+            select_version=select_version,
         )
 
         if model_parameters:
@@ -604,10 +601,8 @@ def create(
             gapic_batch_prediction_job.generate_explanation = generate_explanation
 
         if explanation_metadata or explanation_parameters:
-            gapic_batch_prediction_job.explanation_spec = (
-                gca_explanation_v1beta1.ExplanationSpec(
-                    metadata=explanation_metadata, parameters=explanation_parameters
-                )
+            gapic_batch_prediction_job.explanation_spec = gca_explanation_v1beta1.ExplanationSpec(
+                metadata=explanation_metadata, parameters=explanation_parameters
             )
 
         # TODO (b/174502913): Support private feature once released
@@ -1088,23 +1083,19 @@ def from_local_script(
                 "should be set using aiplatform.init(staging_bucket='gs://my-bucket')"
             )
 
-        worker_pool_specs = (
-            worker_spec_utils._DistributedTrainingSpec.chief_worker_pool(
-                replica_count=replica_count,
-                machine_type=machine_type,
-                accelerator_count=accelerator_count,
-                accelerator_type=accelerator_type,
-            ).pool_specs
-        )
+        worker_pool_specs = worker_spec_utils._DistributedTrainingSpec.chief_worker_pool(
+            replica_count=replica_count,
+            machine_type=machine_type,
+            accelerator_count=accelerator_count,
+            accelerator_type=accelerator_type,
+        ).pool_specs
 
         python_packager = source_utils._TrainingScriptPythonPackager(
             script_path=script_path, requirements=requirements
         )
 
         package_gcs_uri = python_packager.package_and_copy_to_gcs(
-            gcs_staging_dir=staging_bucket,
-            project=project,
-            credentials=credentials,
+            gcs_staging_dir=staging_bucket, project=project, credentials=credentials,
         )
 
         for spec in worker_pool_specs:
@@ -1429,18 +1420,16 @@ def __init__(
             ],
         )
 
-        self._gca_resource = (
-            gca_hyperparameter_tuning_job_compat.HyperparameterTuningJob(
-                display_name=display_name,
-                study_spec=study_spec,
-                max_trial_count=max_trial_count,
-                parallel_trial_count=parallel_trial_count,
-                max_failed_trial_count=max_failed_trial_count,
-                trial_job_spec=copy.deepcopy(custom_job.job_spec),
-                encryption_spec=initializer.global_config.get_encryption_spec(
-                    encryption_spec_key_name=encryption_spec_key_name
-                ),
-            )
+        self._gca_resource = gca_hyperparameter_tuning_job_compat.HyperparameterTuningJob(
+            display_name=display_name,
+            study_spec=study_spec,
+            max_trial_count=max_trial_count,
+            parallel_trial_count=parallel_trial_count,
+            max_failed_trial_count=max_failed_trial_count,
+            trial_job_spec=copy.deepcopy(custom_job.job_spec),
+            encryption_spec=initializer.global_config.get_encryption_spec(
+                encryption_spec_key_name=encryption_spec_key_name
+            ),
         )
 
     @base.optional_sync()
@@ -1499,11 +1488,9 @@ def run(
 
         if timeout or restart_job_on_worker_restart:
             duration = duration_pb2.Duration(seconds=timeout) if timeout else None
-            self._gca_resource.trial_job_spec.scheduling = (
-                gca_custom_job_compat.Scheduling(
-                    timeout=duration,
-                    restart_job_on_worker_restart=restart_job_on_worker_restart,
-                )
+            self._gca_resource.trial_job_spec.scheduling = gca_custom_job_compat.Scheduling(
+                timeout=duration,
+                restart_job_on_worker_restart=restart_job_on_worker_restart,
             )
 
         if tensorboard:

google/cloud/aiplatform/models.py

@@ -297,8 +297,7 @@ def _create(
 
     @staticmethod
     def _allocate_traffic(
-        traffic_split: Dict[str, int],
-        traffic_percentage: int,
+        traffic_split: Dict[str, int], traffic_percentage: int,
     ) -> Dict[str, int]:
         """Allocates desired traffic to new deployed model and scales traffic
         of older deployed models.
@@ -334,8 +333,7 @@ def _allocate_traffic(
 
     @staticmethod
     def _unallocate_traffic(
-        traffic_split: Dict[str, int],
-        deployed_model_id: str,
+        traffic_split: Dict[str, int], deployed_model_id: str,
     ) -> Dict[str, int]:
         """Sets deployed model id's traffic to 0 and scales the traffic of
         other deployed models.
@@ -829,20 +827,16 @@ def _deploy_call(
                 machine_spec.accelerator_type = accelerator_type
                 machine_spec.accelerator_count = accelerator_count
 
-            deployed_model.dedicated_resources = (
-                gca_machine_resources.DedicatedResources(
-                    machine_spec=machine_spec,
-                    min_replica_count=min_replica_count,
-                    max_replica_count=max_replica_count,
-                )
+            deployed_model.dedicated_resources = gca_machine_resources.DedicatedResources(
+                machine_spec=machine_spec,
+                min_replica_count=min_replica_count,
+                max_replica_count=max_replica_count,
             )
 
         else:
-            deployed_model.automatic_resources = (
-                gca_machine_resources.AutomaticResources(
-                    min_replica_count=min_replica_count,
-                    max_replica_count=max_replica_count,
-                )
+            deployed_model.automatic_resources = gca_machine_resources.AutomaticResources(
+                min_replica_count=min_replica_count,
+                max_replica_count=max_replica_count,
             )
 
         # Service will throw error if both metadata and parameters are not provided
@@ -2178,8 +2172,8 @@ def export_model(
             )
 
         if image_destination:
-            output_config.image_destination = (
-                gca_io_compat.ContainerRegistryDestination(output_uri=image_destination)
+            output_config.image_destination = gca_io_compat.ContainerRegistryDestination(
+                output_uri=image_destination
             )
 
         _LOGGER.log_action_start_against_resource("Exporting", "model", self)

google/cloud/aiplatform/training_jobs.py

@@ -616,9 +616,7 @@ def _get_model(self) -> Optional[models.Model]:
             )
 
             return models.Model(
-                fields.id,
-                project=fields.project,
-                location=fields.location,
+                fields.id, project=fields.project, location=fields.location,
             )
 
     def _block_until_complete(self):
@@ -1056,14 +1054,12 @@ def _prepare_and_validate_run(
             model_display_name = model_display_name or self._display_name + "-model"
 
         # validates args and will raise
-        worker_pool_specs = (
-            worker_spec_utils._DistributedTrainingSpec.chief_worker_pool(
-                replica_count=replica_count,
-                machine_type=machine_type,
-                accelerator_count=accelerator_count,
-                accelerator_type=accelerator_type,
-            ).pool_specs
-        )
+        worker_pool_specs = worker_spec_utils._DistributedTrainingSpec.chief_worker_pool(
+            replica_count=replica_count,
+            machine_type=machine_type,
+            accelerator_count=accelerator_count,
+            accelerator_type=accelerator_type,
+        ).pool_specs
 
         managed_model = self._managed_model
         if model_display_name:
@@ -4442,10 +4438,8 @@ def __init__(
                 schema.training_job.definition.automl_text_classification
             )
 
-            training_task_inputs_dict = (
-                training_job_inputs.AutoMlTextClassificationInputs(
-                    multi_label=multi_label
-                )
+            training_task_inputs_dict = training_job_inputs.AutoMlTextClassificationInputs(
+                multi_label=multi_label
             )
         elif prediction_type == "extraction":
             training_task_definition = (

google/cloud/aiplatform/utils/json_utils.py

@@ -22,10 +22,11 @@
 from google.cloud import storage
 
 
-def load_json(path: str,
-              project: Optional[str] = None,
-              credentials: Optional[auth_credentials.Credentials] = None
-              ) -> Dict[str, Any]:
+def load_json(
+    path: str,
+    project: Optional[str] = None,
+    credentials: Optional[auth_credentials.Credentials] = None,
+) -> Dict[str, Any]:
     """Loads data from a JSON document.
 
     Args:
@@ -40,16 +41,17 @@ def load_json(path: str,
     Returns:
       A Dict object representing the JSON document.
     """
-    if path.startswith('gs://'):
+    if path.startswith("gs://"):
         return _load_json_from_gs_uri(path, project, credentials)
     else:
         return _load_json_from_local_file(path)
 
 
-def _load_json_from_gs_uri(uri: str,
-                           project: Optional[str] = None,
-                           credentials: Optional[auth_credentials.Credentials]
-                           = None) -> Dict[str, Any]:
+def _load_json_from_gs_uri(
+    uri: str,
+    project: Optional[str] = None,
+    credentials: Optional[auth_credentials.Credentials] = None,
+) -> Dict[str, Any]:
     """Loads data from a JSON document referenced by a GCS URI.
 
     Args: