fix: CustomJob.from_local_script does not pass args to script for cus…

…tom container images PiperOrigin-RevId: 531282139
googleapis · May 11, 2023 · 6ead69d · 6ead69d
1 parent ad7ef73
commit 6ead69d
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/google/cloud/aiplatform/jobs.py b/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}

diff --git a/tests/unit/aiplatform/test_custom_job.py b/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()