Fixed pipeline step to dict.

Author: lu-ohai

ads/pipeline/ads_pipeline_step.py

@@ -519,6 +519,19 @@ def to_dict(self) -> dict:
             dict_details["spec"][self.CONST_JOB_ID] = self.job_id
         if self.description:
             dict_details["spec"][self.CONST_DESCRIPTION] = self.description
+        if self.kind == "ML_JOB":
+            if self.environment_variable:
+                dict_details["spec"][self.CONST_ENVIRONMENT_VARIABLES] = (
+                    self.environment_variable
+                )
+            if self.argument:
+                dict_details["spec"][self.CONST_COMMAND_LINE_ARGUMENTS] = (
+                    self.argument
+                )
+            if self.maximum_runtime_in_minutes:
+                dict_details["spec"][self.CONST_MAXIMUM_RUNTIME_IN_MINUTES] = (
+                    self.maximum_runtime_in_minutes
+                )
 
         dict_details["spec"].pop(self.CONST_DEPENDS_ON, None)
 

tests/unitary/default_setup/pipeline/test_pipeline_step.py

@@ -38,6 +38,9 @@ class DataSciencePipelineStepBaseTest(unittest.TestCase):
         PipelineStep("TestUpstreamPipelineStepOne")
         .with_description("Test upstream pipeline step description one")
         .with_job_id("TestJobIdOne")
+        .with_environment_variable(**{"test_key": "test_value"})
+        .with_maximum_runtime_in_minutes(20)
+        .with_argument("--key val path/to/file")
     )
 
     upstream_pipeline_step_two = (
@@ -85,6 +88,11 @@ def test_pipeline_step_to_dict(self):
                 "name": "TestUpstreamPipelineStepOne",
                 "jobId": "TestJobIdOne",
                 "description": "Test upstream pipeline step description one",
+                "environmentVariables": {
+                    "test_key": "test_value"
+                },
+                "commandLineArguments": "--key val path/to/file",
+                "maximumRuntimeInMinutes": 20
             },
         }