fix: CustomJob.from_local_script does not pass args to script for custom container images

PiperOrigin-RevId: 531282139

Author: jaycee-li

google/cloud/aiplatform/jobs.py

@@ -1577,14 +1577,14 @@ def from_local_script(
                     + f"python3 -m {python_packager.module_name}",
                 ]
 
+                if args:
+                    command[-1] += " " + " ".join(args)
+
                 spec["container_spec"] = {
                     "image_uri": container_uri,
                     "command": command,
                 }
 
-                if args:
-                    spec["container_spec"]["args"] = args
-
                 if environment_variables:
                     spec["container_spec"]["env"] = [
                         {"name": key, "value": value}

tests/unit/aiplatform/test_custom_job.py

@@ -937,8 +937,8 @@ def test_create_from_local_script_custom_container_with_all_args(
             enable_autolog=True,
         )
 
-        expected_container_spec = _TEST_CONTAINER_SPEC
-        expected_container_spec.args = _TEST_RUN_ARGS
+        expected_container_spec = copy.deepcopy(_TEST_CONTAINER_SPEC)
+        expected_container_spec.command[-1] += " " + " ".join(_TEST_RUN_ARGS)
         expected_container_spec.env = [
             {"name": key, "value": value}
             for key, value in test_constants.TrainingJobConstants._TEST_ENVIRONMENT_VARIABLES.items()