fix fixture issue

Author: TheMichaelHu

google/cloud/aiplatform/training_jobs.py

@@ -1562,12 +1562,10 @@ def _model_upload_fail_string(self) -> str:
 
 
 class _ForecastingTrainingJob(_TrainingJob):
- """ABC for Forecasting Training Pipelines.."""
+ """ABC for Forecasting Training Pipelines."""
 
  def __init__(
  self,
- model_type: str,
- training_task_definition: str,
  display_name: Optional[str] = None,
  optimization_objective: Optional[str] = None,
  column_specs: Optional[Dict[str, str]] = None,
@@ -1582,15 +1580,6 @@ def __init__(
  """Constructs a Forecasting Training Job.
 
  Args:
- model_type (str): The type of forecasting model.
- training_task_definition (str):
- Required. A Google Cloud Storage path to the
- YAML file that defines the training task which
- is responsible for producing the model artifact,
- and may also include additional auxiliary work.
- The definition files that can be used here are
- found in gs://google-cloud-
- aiplatform/schema/trainingjob/definition/.
  display_name (str):
  Optional. The user-defined name of this TrainingPipeline.
  optimization_objective (str):
@@ -1687,8 +1676,24 @@ def __init__(
 
  self._optimization_objective = optimization_objective
  self._additional_experiments = []
- self._model_type = model_type
- self._training_task_definition = training_task_definition
+
+ @property
+ @classmethod
+ @abc.abstractmethod
+ def _model_type(cls) -> str:
+ """The type of forecasting model."""
+ pass
+
+ @property
+ @classmethod
+ @abc.abstractmethod
+ def _training_task_definition(cls) -> str:
+ """A GCS path to the YAML file that defines the training task.
+
+ The definition files that can be used here are found in
+ gs://google-cloud- aiplatform/schema/trainingjob/definition/.
+ """
+ pass
 
  def run(
  self,
@@ -1900,13 +1905,14 @@ def run(
 
  if self._is_waiting_to_run():
  raise RuntimeError(
- f"{self._model_type} Forecasting Training is already scheduled "
- "to run."
+ f"{self.__class__._model_type} Forecasting Training is already "
+ "scheduled to run."
  )
 
  if self._has_run:
  raise RuntimeError(
- f"{self._model_type} Forecasting Training has " "already run."
+ f"{self.__class__._model_type} Forecasting Training has "
+ "already run."
  )
 
  if additional_experiments:
@@ -2210,7 +2216,7 @@ def _run(
  )
 
  new_model = self._run_job(
- training_task_definition=self._training_task_definition,
+ training_task_definition=self.__class__._training_task_definition,
  training_task_inputs=training_task_inputs_dict,
  dataset=dataset,
  training_fraction_split=training_fraction_split,
@@ -4590,6 +4596,9 @@ class column_data_types:
 
 
 class AutoMLForecastingTrainingJob(_ForecastingTrainingJob):
+ _model_type = "AutoML"
+ _training_task_definition = (
+ schema.training_job.definition.automl_forecasting)
  _supported_training_schemas = (schema.training_job.definition.automl_forecasting,)
 
  def __init__(
@@ -4692,10 +4701,6 @@ def __init__(
  ValueError: If both column_transformations and column_specs were provided.
  """
  super().__init__(
- model_type="AutoML",
- training_task_definition=(
- schema.training_job.definition.automl_forecasting
- ),
  display_name=display_name,
  optimization_objective=optimization_objective,
  column_specs=column_specs,
@@ -4954,6 +4959,9 @@ def evaluated_data_items_bigquery_uri(self) -> Optional[str]:
 
 
 class SequenceToSequencePlusForecastingTrainingJob(_ForecastingTrainingJob):
+ _model_type = "Seq2Seq"
+ _training_task_definition = (
+ schema.training_job.definition.seq2seq_forecasting)
  _supported_training_schemas = (schema.training_job.definition.seq2seq_forecasting,)
 
  def __init__(
@@ -5056,10 +5064,6 @@ def __init__(
  ValueError: If both column_transformations and column_specs were provided.
  """
  super().__init__(
- model_type="Seq2Seq",
- training_task_definition=(
- schema.training_job.definition.seq2seq_forecasting
- ),
  display_name=display_name,
  optimization_objective=optimization_objective,
  column_specs=column_specs,

tests/unit/aiplatform/test_automl_forecasting_training_jobs.py

@@ -256,9 +256,8 @@ def setup_method(self):
  def teardown_method(self):
  initializer.global_pool.shutdown(wait=True)
 
+ @pytest.mark.parametrize("sync", [True, False])
  @pytest.mark.parametrize(
- "sync",
- [True, False],
  "training_job",
  [
  training_jobs.AutoMLForecastingTrainingJob,
@@ -327,7 +326,7 @@ def test_run_call_pipeline_service_create(
  true_training_pipeline = gca_training_pipeline.TrainingPipeline(
  display_name=_TEST_DISPLAY_NAME,
  labels=_TEST_LABELS,
- training_task_definition=schema.training_job.definition.automl_forecasting,
+ training_task_definition=training_job._training_task_definition,
  training_task_inputs=_TEST_TRAINING_TASK_INPUTS_WITH_ADDITIONAL_EXPERIMENTS,
  model_to_upload=true_managed_model,
  input_data_config=true_input_data_config,
@@ -353,9 +352,8 @@ def test_run_call_pipeline_service_create(
 
  assert job.state == gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED
 
+ @pytest.mark.parametrize("sync", [True, False])
  @pytest.mark.parametrize(
- "sync",
- [True, False],
  "training_job",
  [
  training_jobs.AutoMLForecastingTrainingJob,
@@ -424,7 +422,7 @@ def test_run_call_pipeline_service_create_with_timeout(
  true_training_pipeline = gca_training_pipeline.TrainingPipeline(
  display_name=_TEST_DISPLAY_NAME,
  labels=_TEST_LABELS,
- training_task_definition=schema.training_job.definition.automl_forecasting,
+ training_task_definition=training_job._training_task_definition,
  training_task_inputs=_TEST_TRAINING_TASK_INPUTS_WITH_ADDITIONAL_EXPERIMENTS,
  model_to_upload=true_managed_model,
  input_data_config=true_input_data_config,
@@ -437,9 +435,8 @@ def test_run_call_pipeline_service_create_with_timeout(
  )
 
  @pytest.mark.usefixtures("mock_pipeline_service_get")
+ @pytest.mark.parametrize("sync", [True, False])
  @pytest.mark.parametrize(
- "sync",
- [True, False],
  "training_job",
  [
  training_jobs.AutoMLForecastingTrainingJob,
@@ -502,7 +499,7 @@ def test_run_call_pipeline_if_no_model_display_name_nor_model_labels(
  true_training_pipeline = gca_training_pipeline.TrainingPipeline(
  display_name=_TEST_DISPLAY_NAME,
  labels=_TEST_LABELS,
- training_task_definition=schema.training_job.definition.automl_forecasting,
+ training_task_definition=training_job._training_task_definition,
  training_task_inputs=_TEST_TRAINING_TASK_INPUTS,
  model_to_upload=true_managed_model,
  input_data_config=true_input_data_config,
@@ -515,9 +512,8 @@ def test_run_call_pipeline_if_no_model_display_name_nor_model_labels(
  )
 
  @pytest.mark.usefixtures("mock_pipeline_service_get")
+ @pytest.mark.parametrize("sync", [True, False])
  @pytest.mark.parametrize(
- "sync",
- [True, False],
  "training_job",
  [
  training_jobs.AutoMLForecastingTrainingJob,
@@ -577,7 +573,7 @@ def test_run_call_pipeline_if_set_additional_experiments(
 
  true_training_pipeline = gca_training_pipeline.TrainingPipeline(
  display_name=_TEST_DISPLAY_NAME,
- training_task_definition=schema.training_job.definition.automl_forecasting,
+ training_task_definition=training_job._training_task_definition,
  training_task_inputs=_TEST_TRAINING_TASK_INPUTS_WITH_ADDITIONAL_EXPERIMENTS,
  model_to_upload=true_managed_model,
  input_data_config=true_input_data_config,
@@ -594,9 +590,8 @@ def test_run_call_pipeline_if_set_additional_experiments(
  "mock_pipeline_service_get",
  "mock_model_service_get",
  )
+ @pytest.mark.parametrize("sync", [True, False])
  @pytest.mark.parametrize(
- "sync",
- [True, False],
  "training_job",
  [
  training_jobs.AutoMLForecastingTrainingJob,
@@ -664,9 +659,8 @@ def test_run_called_twice_raises(
  sync=sync,
  )
 
+ @pytest.mark.parametrize("sync", [True, False])
  @pytest.mark.parametrize(
- "sync",
- [True, False],
  "training_job",
  [
  training_jobs.AutoMLForecastingTrainingJob,
@@ -748,9 +742,8 @@ def test_raises_before_run_is_called(
  with pytest.raises(RuntimeError):
  job.state
 
+ @pytest.mark.parametrize("sync", [True, False])
  @pytest.mark.parametrize(
- "sync",
- [True, False],
  "training_job",
  [
  training_jobs.AutoMLForecastingTrainingJob,
@@ -830,7 +823,7 @@ def test_splits_fraction(
 
  true_training_pipeline = gca_training_pipeline.TrainingPipeline(
  display_name=_TEST_DISPLAY_NAME,
- training_task_definition=schema.training_job.definition.automl_forecasting,
+ training_task_definition=training_job._training_task_definition,
  training_task_inputs=_TEST_TRAINING_TASK_INPUTS,
  model_to_upload=true_managed_model,
  input_data_config=true_input_data_config,
@@ -843,9 +836,8 @@ def test_splits_fraction(
  timeout=None,
  )
 
+ @pytest.mark.parametrize("sync", [True, False])
  @pytest.mark.parametrize(
- "sync",
- [True, False],
  "training_job",
  [
  training_jobs.AutoMLForecastingTrainingJob,
@@ -927,9 +919,7 @@ def test_splits_timestamp(
 
  true_training_pipeline = gca_training_pipeline.TrainingPipeline(
  display_name=_TEST_DISPLAY_NAME,
- training_task_definition=(
- schema.training_job.definition.automl_forecasting
- ),
+ training_task_definition=training_job._training_task_definition,
  training_task_inputs=_TEST_TRAINING_TASK_INPUTS,
  model_to_upload=true_managed_model,
  input_data_config=true_input_data_config,
@@ -942,9 +932,8 @@ def test_splits_timestamp(
  timeout=None,
  )
 
+ @pytest.mark.parametrize("sync", [True, False])
  @pytest.mark.parametrize(
- "sync",
- [True, False],
  "training_job",
  [
  training_jobs.AutoMLForecastingTrainingJob,
@@ -1020,7 +1009,7 @@ def test_splits_predefined(
 
  true_training_pipeline = gca_training_pipeline.TrainingPipeline(
  display_name=_TEST_DISPLAY_NAME,
- training_task_definition=schema.training_job.definition.automl_forecasting,
+ training_task_definition=training_job._training_task_definition,
  training_task_inputs=_TEST_TRAINING_TASK_INPUTS,
  model_to_upload=true_managed_model,
  input_data_config=true_input_data_config,
@@ -1033,9 +1022,8 @@ def test_splits_predefined(
  timeout=None,
  )
 
+ @pytest.mark.parametrize("sync", [True, False])
  @pytest.mark.parametrize(
- "sync",
- [True, False],
  "training_job",
  [
  training_jobs.AutoMLForecastingTrainingJob,
@@ -1105,7 +1093,7 @@ def test_splits_default(
 
  true_training_pipeline = gca_training_pipeline.TrainingPipeline(
  display_name=_TEST_DISPLAY_NAME,
- training_task_definition=schema.training_job.definition.automl_forecasting,
+ training_task_definition=training_job._training_task_definition,
  training_task_inputs=_TEST_TRAINING_TASK_INPUTS,
  model_to_upload=true_managed_model,
  input_data_config=true_input_data_config,