Fixed pipeline step to dict (#325)

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.py

@@ -333,6 +333,11 @@ def test_pipeline_define(self):
                             "name": "TestPipelineStepOne",
                             "jobId": "TestJobIdOne",
                             "description": "Test description one",
+                            "commandLineArguments": "ARGUMENT --KEY VALUE",
+                            "environmentVariables": {
+                                "ENV": "VALUE"
+                            },
+                            "maximumRuntimeInMinutes": 20
                         },
                     },
                     {
@@ -1060,6 +1065,11 @@ def test_pipeline_to_dict(self):
                             "name": "TestPipelineStepOne",
                             "jobId": "TestJobIdOne",
                             "description": "Test description one",
+                            "commandLineArguments": "ARGUMENT --KEY VALUE",
+                            "environmentVariables": {
+                                "ENV": "VALUE"
+                            },
+                            "maximumRuntimeInMinutes": 20
                         },
                     },
                     {

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
             },
         }