CustomContainerTrainingJob.run drops max_wait_duration=0 instead of requesting indefinite DWS wait

[someshfengade]

Thanks for maintaining this library. CustomContainerTrainingJob.run(max_wait_duration=0) drops the value because the high-level training-job implementation uses a truthiness check. This prevents callers from requesting the documented indefinite wait for DWS Flex Start.

Environment details

• OS type and version: macOS
• Python version: 3.12
• google-cloud-aiplatform versions: reproduced in 1.148.1; source inspection confirms the same behavior in 1.150.0, 1.163.0, and current main

Steps to reproduce

1. Create a CustomContainerTrainingJob using any valid project, staging bucket, and container image.
2. Run it with Flex Start and max_wait_duration=0.
3. Inspect the scheduling inputs produced by _prepare_training_task_inputs_and_output_dir or the resulting TrainingPipeline request.

Code example

job.run(
    scheduling_strategy=custom_job.Scheduling.Strategy.FLEX_START,
    max_wait_duration=0,
)

Observed scheduling value:

max_wait_duration: None

Expected scheduling value:

max_wait_duration: 0s

For comparison, a positive input such as 14400 is preserved as 14400s.

Cause

The high-level training-job path converts the duration using a truthiness check, so zero follows the same branch as None:

https://github.com/googleapis/python-aiplatform/blob/v1.163.0/google/cloud/aiplatform/training_jobs.py#L1663-L1667

The generated API contract says that explicit zero means indefinite waiting, while omission defaults to 24 hours:

https://github.com/googleapis/python-aiplatform/blob/v1.163.0/google/cloud/aiplatform_v1/types/custom_job.py#L577-L582

max_wait_duration is a presence-aware google.protobuf.Duration, so explicit zero and an omitted field are distinct requests.

A previous fix correctly changed the lower-level CustomJob and hyperparameter-tuning paths to use is not None, explicitly noting that 0 is valid:

d9675fd

That commit did not update the CustomContainerTrainingJob path in training_jobs.py.

Suggested fix

Use an explicit is not None check in _prepare_training_task_inputs_and_output_dir, consistent with the lower-level implementation, and add tests covering omitted, zero, and positive durations.

Stack trace

No exception is raised. The problem is a silent request-semantics change: 0 is serialized as omission.

[someshfengade]

The official CustomContainerTrainingJob.run() reference states that max_wait_duration=0 waits indefinitely and that the omitted default is 30 minutes:

https://docs.cloud.google.com/python/docs/reference/aiplatform/1.148.1/google.cloud.aiplatform.CustomContainerTrainingJob#google_cloud_aiplatform_CustomContainerTrainingJob_run

Final serialized scheduling results

SDK input	Serialized max_wait_duration
None	None
0	None
14400	14400s

Thus, SDK input 0 produces exactly the same final request value as omission.

Protobuf presence control

I also constructed Scheduling directly to confirm that the API type can represent explicit zero and distinguishes it from omission:

Construction	HasField(max_wait_duration)	JSON
omitted	False	strategy only
explicit 0s	True	strategy plus max_wait_duration: 0s

This rules out proto3 default-value behavior as the cause. The value is lost earlier by the truthiness conversion in training_jobs.py; explicit zero would survive serialization if that conversion preserved it.