feat: expose env var in cust training class run func args (#366)

Author: morgandu

google/cloud/aiplatform/training_jobs.py

@@ -1805,6 +1805,7 @@ def run(
  service_account: Optional[str] = None,
  bigquery_destination: Optional[str] = None,
  args: Optional[List[Union[str, float, int]]] = None,
+ environment_variables: Optional[Dict[str, str]] = None,
  replica_count: int = 0,
  machine_type: str = "n1-standard-4",
  accelerator_type: str = "ACCELERATOR_TYPE_UNSPECIFIED",
@@ -1880,6 +1881,13 @@ def run(
  base_output_dir (str):
  GCS output directory of job. If not provided a
  timestamped directory in the staging directory will be used.
+
+ AI Platform sets the following environment variables when it runs your training code:
+
+ - AIP_MODEL_DIR: a Cloud Storage URI of a directory intended for saving model artifacts, i.e. <base_output_dir>/model/
+ - AIP_CHECKPOINT_DIR: a Cloud Storage URI of a directory intended for saving checkpoints, i.e. <base_output_dir>/checkpoints/
+ - AIP_TENSORBOARD_LOG_DIR: a Cloud Storage URI of a directory intended for saving TensorBoard logs, i.e. <base_output_dir>/logs/
+
  service_account (str):
  Specifies the service account for workload run-as account.
  Users submitting jobs must have act-as permission on this run-as account.
@@ -1900,6 +1908,16 @@ def run(
  - AIP_TEST_DATA_URI = "bigquery_destination.dataset_*.test"
  args (List[Unions[str, int, float]]):
  Command line arguments to be passed to the Python script.
+ environment_variables (Dict[str, str]):
+ Environment variables to be passed to the container.
+ Should be a dictionary where keys are environment variable names
+ and values are environment variable values for those names.
+ At most 10 environment variables can be specified.
+ The Name of the environment variable must be unique.
+
+ environment_variables = {
+ 'MY_KEY': 'MY_VALUE'
+ }
  replica_count (int):
  The number of worker replicas. If replica count = 1 then one chief
  replica will be provisioned. If replica_count > 1 the remainder will be
@@ -1960,6 +1978,7 @@ def run(
  worker_pool_specs=worker_pool_specs,
  managed_model=managed_model,
  args=args,
+ environment_variables=environment_variables,
  base_output_dir=base_output_dir,
  service_account=service_account,
  bigquery_destination=bigquery_destination,
@@ -1986,6 +2005,7 @@ def _run(
  worker_pool_specs: _DistributedTrainingSpec,
  managed_model: Optional[gca_model.Model] = None,
  args: Optional[List[Union[str, float, int]]] = None,
+ environment_variables: Optional[Dict[str, str]] = None,
  base_output_dir: Optional[str] = None,
  service_account: Optional[str] = None,
  bigquery_destination: Optional[str] = None,
@@ -2018,9 +2038,26 @@ def _run(
  Model proto if this script produces a Managed Model.
  args (List[Unions[str, int, float]]):
  Command line arguments to be passed to the Python script.
+ environment_variables (Dict[str, str]):
+ Environment variables to be passed to the container.
+ Should be a dictionary where keys are environment variable names
+ and values are environment variable values for those names.
+ At most 10 environment variables can be specified.
+ The Name of the environment variable must be unique.
+
+ environment_variables = {
+ 'MY_KEY': 'MY_VALUE'
+ }
  base_output_dir (str):
  GCS output directory of job. If not provided a
  timestamped directory in the staging directory will be used.
+
+ AI Platform sets the following environment variables when it runs your training code:
+
+ - AIP_MODEL_DIR: a Cloud Storage URI of a directory intended for saving model artifacts, i.e. <base_output_dir>/model/
+ - AIP_CHECKPOINT_DIR: a Cloud Storage URI of a directory intended for saving checkpoints, i.e. <base_output_dir>/checkpoints/
+ - AIP_TENSORBOARD_LOG_DIR: a Cloud Storage URI of a directory intended for saving TensorBoard logs, i.e. <base_output_dir>/logs/
+
  service_account (str):
  Specifies the service account for workload run-as account.
  Users submitting jobs must have act-as permission on this run-as account.
@@ -2083,6 +2120,9 @@ def _run(
  if args:
  spec["pythonPackageSpec"]["args"] = args
 
+ if environment_variables:
+ spec["pythonPackageSpec"]["env"] = environment_variables
+
  (
  training_task_inputs,
  base_output_dir,
@@ -2334,6 +2374,7 @@ def run(
  service_account: Optional[str] = None,
  bigquery_destination: Optional[str] = None,
  args: Optional[List[Union[str, float, int]]] = None,
+ environment_variables: Optional[Dict[str, str]] = None,
  replica_count: int = 0,
  machine_type: str = "n1-standard-4",
  accelerator_type: str = "ACCELERATOR_TYPE_UNSPECIFIED",
@@ -2402,6 +2443,13 @@ def run(
  base_output_dir (str):
  GCS output directory of job. If not provided a
  timestamped directory in the staging directory will be used.
+
+ AI Platform sets the following environment variables when it runs your training code:
+
+ - AIP_MODEL_DIR: a Cloud Storage URI of a directory intended for saving model artifacts, i.e. <base_output_dir>/model/
+ - AIP_CHECKPOINT_DIR: a Cloud Storage URI of a directory intended for saving checkpoints, i.e. <base_output_dir>/checkpoints/
+ - AIP_TENSORBOARD_LOG_DIR: a Cloud Storage URI of a directory intended for saving TensorBoard logs, i.e. <base_output_dir>/logs/
+
  service_account (str):
  Specifies the service account for workload run-as account.
  Users submitting jobs must have act-as permission on this run-as account.
@@ -2422,6 +2470,16 @@ def run(
  - AIP_TEST_DATA_URI = "bigquery_destination.dataset_*.test"
  args (List[Unions[str, int, float]]):
  Command line arguments to be passed to the Python script.
+ environment_variables (Dict[str, str]):
+ Environment variables to be passed to the container.
+ Should be a dictionary where keys are environment variable names
+ and values are environment variable values for those names.
+ At most 10 environment variables can be specified.
+ The Name of the environment variable must be unique.
+
+ environment_variables = {
+ 'MY_KEY': 'MY_VALUE'
+ }
  replica_count (int):
  The number of worker replicas. If replica count = 1 then one chief
  replica will be provisioned. If replica_count > 1 the remainder will be
@@ -2481,6 +2539,7 @@ def run(
  worker_pool_specs=worker_pool_specs,
  managed_model=managed_model,
  args=args,
+ environment_variables=environment_variables,
  base_output_dir=base_output_dir,
  service_account=service_account,
  bigquery_destination=bigquery_destination,
@@ -2506,6 +2565,7 @@ def _run(
  worker_pool_specs: _DistributedTrainingSpec,
  managed_model: Optional[gca_model.Model] = None,
  args: Optional[List[Union[str, float, int]]] = None,
+ environment_variables: Optional[Dict[str, str]] = None,
  base_output_dir: Optional[str] = None,
  service_account: Optional[str] = None,
  bigquery_destination: Optional[str] = None,
@@ -2535,9 +2595,26 @@ def _run(
  Model proto if this script produces a Managed Model.
  args (List[Unions[str, int, float]]):
  Command line arguments to be passed to the Python script.
+ environment_variables (Dict[str, str]):
+ Environment variables to be passed to the container.
+ Should be a dictionary where keys are environment variable names
+ and values are environment variable values for those names.
+ At most 10 environment variables can be specified.
+ The Name of the environment variable must be unique.
+
+ environment_variables = {
+ 'MY_KEY': 'MY_VALUE'
+ }
  base_output_dir (str):
  GCS output directory of job. If not provided a
  timestamped directory in the staging directory will be used.
+
+ AI Platform sets the following environment variables when it runs your training code:
+
+ - AIP_MODEL_DIR: a Cloud Storage URI of a directory intended for saving model artifacts, i.e. <base_output_dir>/model/
+ - AIP_CHECKPOINT_DIR: a Cloud Storage URI of a directory intended for saving checkpoints, i.e. <base_output_dir>/checkpoints/
+ - AIP_TENSORBOARD_LOG_DIR: a Cloud Storage URI of a directory intended for saving TensorBoard logs, i.e. <base_output_dir>/logs/
+
  service_account (str):
  Specifies the service account for workload run-as account.
  Users submitting jobs must have act-as permission on this run-as account.
@@ -2593,6 +2670,9 @@ def _run(
  if args:
  spec["containerSpec"]["args"] = args
 
+ if environment_variables:
+ spec["containerSpec"]["env"] = environment_variables
+
  (
  training_task_inputs,
  base_output_dir,
@@ -3625,6 +3705,7 @@ def run(
  service_account: Optional[str] = None,
  bigquery_destination: Optional[str] = None,
  args: Optional[List[Union[str, float, int]]] = None,
+ environment_variables: Optional[Dict[str, str]] = None,
  replica_count: int = 0,
  machine_type: str = "n1-standard-4",
  accelerator_type: str = "ACCELERATOR_TYPE_UNSPECIFIED",
@@ -3693,6 +3774,13 @@ def run(
  base_output_dir (str):
  GCS output directory of job. If not provided a
  timestamped directory in the staging directory will be used.
+
+ AI Platform sets the following environment variables when it runs your training code:
+
+ - AIP_MODEL_DIR: a Cloud Storage URI of a directory intended for saving model artifacts, i.e. <base_output_dir>/model/
+ - AIP_CHECKPOINT_DIR: a Cloud Storage URI of a directory intended for saving checkpoints, i.e. <base_output_dir>/checkpoints/
+ - AIP_TENSORBOARD_LOG_DIR: a Cloud Storage URI of a directory intended for saving TensorBoard logs, i.e. <base_output_dir>/logs/
+
  service_account (str):
  Specifies the service account for workload run-as account.
  Users submitting jobs must have act-as permission on this run-as account.
@@ -3713,6 +3801,16 @@ def run(
  - AIP_TEST_DATA_URI = "bigquery_destination.dataset_*.test"
  args (List[Unions[str, int, float]]):
  Command line arguments to be passed to the Python script.
+ environment_variables (Dict[str, str]):
+ Environment variables to be passed to the container.
+ Should be a dictionary where keys are environment variable names
+ and values are environment variable values for those names.
+ At most 10 environment variables can be specified.
+ The Name of the environment variable must be unique.
+
+ environment_variables = {
+ 'MY_KEY': 'MY_VALUE'
+ }
  replica_count (int):
  The number of worker replicas. If replica count = 1 then one chief
  replica will be provisioned. If replica_count > 1 the remainder will be
@@ -3767,6 +3865,7 @@ def run(
  worker_pool_specs=worker_pool_specs,
  managed_model=managed_model,
  args=args,
+ environment_variables=environment_variables,
  base_output_dir=base_output_dir,
  service_account=service_account,
  training_fraction_split=training_fraction_split,
@@ -3792,6 +3891,7 @@ def _run(
  worker_pool_specs: _DistributedTrainingSpec,
  managed_model: Optional[gca_model.Model] = None,
  args: Optional[List[Union[str, float, int]]] = None,
+ environment_variables: Optional[Dict[str, str]] = None,
  base_output_dir: Optional[str] = None,
  service_account: Optional[str] = None,
  training_fraction_split: float = 0.8,
@@ -3822,9 +3922,26 @@ def _run(
  Model proto if this script produces a Managed Model.
  args (List[Unions[str, int, float]]):
  Command line arguments to be passed to the Python script.
+ environment_variables (Dict[str, str]):
+ Environment variables to be passed to the container.
+ Should be a dictionary where keys are environment variable names
+ and values are environment variable values for those names.
+ At most 10 environment variables can be specified.
+ The Name of the environment variable must be unique.
+
+ environment_variables = {
+ 'MY_KEY': 'MY_VALUE'
+ }
  base_output_dir (str):
  GCS output directory of job. If not provided a
  timestamped directory in the staging directory will be used.
+
+ AI Platform sets the following environment variables when it runs your training code:
+
+ - AIP_MODEL_DIR: a Cloud Storage URI of a directory intended for saving model artifacts, i.e. <base_output_dir>/model/
+ - AIP_CHECKPOINT_DIR: a Cloud Storage URI of a directory intended for saving checkpoints, i.e. <base_output_dir>/checkpoints/
+ - AIP_TENSORBOARD_LOG_DIR: a Cloud Storage URI of a directory intended for saving TensorBoard logs, i.e. <base_output_dir>/logs/
+
  service_account (str):
  Specifies the service account for workload run-as account.
  Users submitting jobs must have act-as permission on this run-as account.
@@ -3866,6 +3983,9 @@ def _run(
  if args:
  spec["pythonPackageSpec"]["args"] = args
 
+ if environment_variables:
+ spec["pythonPackageSpec"]["env"] = environment_variables
+
  (
  training_task_inputs,
  base_output_dir,