[Python] Limit custom job duration for vertex AI pipeline

For testing purposes I need to limit a custom job execution to a low value, 1-60 seconds. timeout parameter according to the documentation is:

timeout (int):
Optional.The maximum duration that a pipeline training job can
run. `timeout` is specified in seconds. The default is 80,400
seconds (7 days).


But seems to not be working, since in google cloud console I see the train step for more that 600 seconds. What would be the best way to limit the execution of the code?

Current code ( example)

from google.cloud import aiplatform
job = aiplatform.CustomTrainingJob(

display_name=job_name,
script_path="task.py",
container_uri=TRAIN_IMAGE,
requirements=["tensorflow_datasets"],
model_serving_container_image_uri=DEPLOY_IMAGE,
)
CMDARGS = [
"--epochs=" + str(epochs),
"--dataset="+str(DATASET),
"--steps=" + str(steps),
"--distribute=" + TRAIN_STRATEGY
]
# Start the training
model = job.run(

model_display_name=MODEL_DISPLAY_NAME,
args=CMDARGS,
timeout=3,
replica_count=1,
)


task.py is a sample code on the cifar dataset using tensorflow

Continue reading...