[ML] Restrict JobServiceType to jupyter_lab, ssh, tensor_board, vs_co…

…de and perform required conversion before sending to rest API (Azure#26863) * [ML] Validate the values of JobService.job_service_type passed to Command() and use snake_case * Fix types list * Fix tests * Update changelog.md
azure-sdk · Oct 25, 2022 · defaa79 · defaa79
1 parent 3871bfa
commit defaa79
Show file tree

Hide file tree

Showing 14 changed files with 119 additions and 49 deletions.
diff --git a/sdk/ml/azure-ai-ml/CHANGELOG.md b/sdk/ml/azure-ai-ml/CHANGELOG.md
@@ -6,6 +6,7 @@
 - Registry list operation now accepts scope value to allow subscription-only based requests.
 - Most configuration classes from the entity package now implement the standard mapping protocol.
 - Add registry delete operation.
+- The values of JobService.job_service_type are now using the snake case. e.g jupyter_lab, ssh, tensor_board, vs_code.
 
 ### Breaking Changes
 

diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/constants/_job/job.py b/sdk/ml/azure-ai-ml/azure/ai/ml/constants/_job/job.py
@@ -72,3 +72,37 @@ class RestSparkConfKey:
     DYNAMIC_ALLOCATION_MIN_EXECUTORS = "spark.dynamicAllocation.minExecutors"
     DYNAMIC_ALLOCATION_MAX_EXECUTORS = "spark.dynamicAllocation.maxExecutors"
     DYNAMIC_ALLOCATION_ENABLED = "spark.dynamicAllocation.enabled"
+
+
+class JobServiceTypeNames:
+    class EntityNames:
+        CUSTOM = "custom"
+        TRACKING = "tracking"
+        STUDIO = "studio"
+        JUPYTER_LAB = "jupyter_lab"
+        SSH = "ssh"
+        TENSOR_BOARD = "tensor_board"
+        VS_CODE = "vs_code"
+
+    class RestNames:
+        CUSTOM = "Custom"
+        TRACKING = "Tracking"
+        STUDIO = "Studio"
+        JUPYTER_LAB = "JupyterLab"
+        SSH = "SSH"
+        TENSOR_BOARD = "TensorBoard"
+        VS_CODE = "VSCode"
+
+    ENTITY_TO_REST = {
+        EntityNames.CUSTOM: RestNames.CUSTOM,
+        EntityNames.TRACKING: RestNames.TRACKING,
+        EntityNames.STUDIO: RestNames.STUDIO,
+        EntityNames.JUPYTER_LAB: RestNames.JUPYTER_LAB,
+        EntityNames.SSH: RestNames.SSH,
+        EntityNames.TENSOR_BOARD: RestNames.TENSOR_BOARD,
+        EntityNames.VS_CODE: RestNames.VS_CODE,
+    }
+
+    REST_TO_ENTITY = {v: k for k, v in ENTITY_TO_REST.items()}
+
+    NAMES_ALLOWED_FOR_PUBLIC = [EntityNames.JUPYTER_LAB, EntityNames.SSH, EntityNames.TENSOR_BOARD, EntityNames.VS_CODE]
diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/job_service.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/job_service.py
@@ -8,6 +8,7 @@
 
 from azure.ai.ml._restclient.v2022_10_01_preview.models import AllNodes
 from azure.ai.ml._restclient.v2022_10_01_preview.models import JobService as RestJobService
+from azure.ai.ml.constants._job.job import JobServiceTypeNames
 from azure.ai.ml.entities._mixins import RestTranslatableMixin
 from azure.ai.ml.exceptions import ErrorCategory, ErrorTarget, ValidationErrorType, ValidationException
 
@@ -37,7 +38,7 @@ def __init__(
         self,
         *,
         endpoint: Optional[str] = None,
-        job_service_type: Optional[Literal["JupyterLab", "SSH", "TensorBoard", "VSCode"]] = None,
+        job_service_type: Optional[Literal["jupyter_lab", "ssh", "tensor_board", "vs_code"]] = None,
         nodes: Optional[Literal["all"]] = None,
         status: Optional[str] = None,
         port: Optional[int] = None,
@@ -51,11 +52,14 @@ def __init__(
         self.port = port
         self.properties = properties
         self._validate_nodes()
+        self._validate_job_service_type_name()
 
     def _to_rest_object(self) -> RestJobService:
         return RestJobService(
             endpoint=self.endpoint,
-            job_service_type=self.job_service_type,
+            job_service_type=JobServiceTypeNames.ENTITY_TO_REST.get(self.job_service_type, None)
+            if self.job_service_type
+            else None,
             nodes=AllNodes() if self.nodes else None,
             status=self.status,
             port=self.port,
@@ -73,6 +77,20 @@ def _validate_nodes(self):
                 error_type=ValidationErrorType.INVALID_VALUE,
             )
 
+    def _validate_job_service_type_name(self):
+        if self.job_service_type and not self.job_service_type in JobServiceTypeNames.ENTITY_TO_REST.keys():
+            msg = (
+                f"job_service_type should be one of "
+                f"{JobServiceTypeNames.NAMES_ALLOWED_FOR_PUBLIC}, but received '{self.job_service_type}'."
+            )
+            raise ValidationException(
+                message=msg,
+                no_personal_data_message=msg,
+                target=ErrorTarget.JOB,
+                error_category=ErrorCategory.USER_ERROR,
+                error_type=ValidationErrorType.INVALID_VALUE,
+            )
+
     @classmethod
     def _to_rest_job_services(cls, services: Dict[str, "JobService"]) -> Dict[str, RestJobService]:
         if services is None:
@@ -85,7 +103,10 @@ def _to_rest_job_services(cls, services: Dict[str, "JobService"]) -> Dict[str, R
     def _from_rest_object(cls, obj: RestJobService) -> "JobService":
         return cls(
             endpoint=obj.endpoint,
-            job_service_type=obj.job_service_type,
+            job_service_type=JobServiceTypeNames.REST_TO_ENTITY.get(obj.job_service_type, None)
+            if obj.job_service_type
+            else None,
+            # obj.job_service_type,
             # nodes="all" if isinstance(obj.nodes, AllNodes) else None,
             nodes="all" if obj.nodes else None,
             status=obj.status,

diff --git a/sdk/ml/azure-ai-ml/tests/dsl/unittests/test_command_builder.py b/sdk/ml/azure-ai-ml/tests/dsl/unittests/test_command_builder.py
@@ -842,9 +842,9 @@ def test_spark_job_with_additional_conf(self):
 
     def test_command_services_nodes(self) -> None:
         services = {
-            "my_jupyterlab": {"job_service_type": "JupyterLab", "nodes": "all"},
+            "my_jupyterlab": {"job_service_type": "jupyter_lab", "nodes": "all"},
             "my_tensorboard": {
-                "job_service_type": "TensorBoard",
+                "job_service_type": "tensor_board",
                 "properties": {
                     "logDir": "~/tblog",
                 },
@@ -877,7 +877,17 @@ def test_command_services_nodes(self) -> None:
 
     def test_command_services(self) -> None:
         services = {
-            "my_jupyter": {"job_service_type": "Jupyter"},
+            "my_ssh": {"job_service_type": "ssh"},
+            "my_tensorboard": {
+                "job_service_type": "tensor_board",
+                "properties": {
+                    "logDir": "~/tblog",
+                },
+            },
+            "my_jupyterlab": {"job_service_type": "jupyter_lab"},
+        }
+        rest_services = {
+            "my_ssh": {"job_service_type": "SSH"},
             "my_tensorboard": {
                 "job_service_type": "TensorBoard",
                 "properties": {
@@ -903,10 +913,10 @@ def test_command_services(self) -> None:
             assert isinstance(service, JobService)
 
         node_rest_obj = node._to_rest_object()
-        assert node_rest_obj["services"] == services
+        assert node_rest_obj["services"] == rest_services
 
         # test invalid services
-        invalid_services_0 = "jupyter"
+        invalid_services_0 = "ssh"
         with pytest.raises(ValidationException, match="Services must be a dict"):
             node = command(
                 name="interactive-command-job",

diff --git a/sdk/ml/azure-ai-ml/tests/dsl/unittests/test_dsl_pipeline_with_specific_nodes.py b/sdk/ml/azure-ai-ml/tests/dsl/unittests/test_dsl_pipeline_with_specific_nodes.py
@@ -1405,6 +1405,7 @@ def pipeline(number, path):
                 "tags": {},
             }
         }
+
     def test_multi_parallel_components_with_file_input_pipeline_output(self) -> None:
         components_dir = tests_root_dir / "test_configs/dsl_pipeline/parallel_component_with_file_input"
         batch_inference1 = load_component(source=str(components_dir / "score.yml"))
@@ -1835,7 +1836,17 @@ def train_with_automl_in_pipeline(training_data, target_column_name_input: str):
 
     def test_pipeline_with_command_services(self):
         services = {
-            "my_jupyter": {"job_service_type": "Jupyter"},
+            "my_ssh": {"job_service_type": "ssh"},
+            "my_tensorboard": {
+                "job_service_type": "tensor_board",
+                "properties": {
+                    "logDir": "~/tblog",
+                },
+            },
+            "my_jupyterlab": {"job_service_type": "jupyter_lab"},
+        }
+        rest_services = {
+            "my_ssh": {"job_service_type": "SSH"},
             "my_tensorboard": {
                 "job_service_type": "TensorBoard",
                 "properties": {
@@ -1874,7 +1885,7 @@ def sample_pipeline():
             assert isinstance(service, JobService)
 
         job_rest_obj = pipeline._to_rest_object()
-        assert job_rest_obj.properties.jobs["node"]["services"] == services
+        assert job_rest_obj.properties.jobs["node"]["services"] == rest_services
 
         recovered_obj = PipelineJob._from_rest_object(job_rest_obj)
         node_services = recovered_obj.jobs["node"].services
@@ -1884,7 +1895,8 @@ def sample_pipeline():
             assert isinstance(service, JobService)
 
         # test set services in pipeline
-        new_services = {"my_jupyter": {"job_service_type": "Jupyter"}}
+        new_services = {"my_jupyter": {"job_service_type": "jupyter_lab"}}
+        rest_new_services = {"my_jupyter": {"job_service_type": "JupyterLab"}}
 
         @dsl.pipeline()
         def sample_pipeline_with_new_services():
@@ -1899,7 +1911,7 @@ def sample_pipeline_with_new_services():
             assert isinstance(service, JobService)
 
         job_rest_obj = pipeline._to_rest_object()
-        assert job_rest_obj.properties.jobs["node"]["services"] == new_services
+        assert job_rest_obj.properties.jobs["node"]["services"] == rest_new_services
 
     def test_pipeline_with_pipeline_component_entity(self):
         path = "./tests/test_configs/components/helloworld_component.yml"

diff --git a/sdk/ml/azure-ai-ml/tests/pipeline_job/unittests/test_pipeline_job_schema.py b/sdk/ml/azure-ai-ml/tests/pipeline_job/unittests/test_pipeline_job_schema.py
@@ -1686,8 +1686,8 @@ def test_command_job_node_services_in_pipeline(self):
         rest_services = job_rest_obj.properties.jobs["hello_world_component_inline"]["services"]
         # rest object of node in pipeline should be pure dict
         assert rest_services == {
-            "my_jupyter": {
-                "job_service_type": "Jupyter",
+            "my_ssh": {
+                "job_service_type": "SSH",
             },
             "my_tensorboard": {
                 "job_service_type": "TensorBoard",
@@ -1710,8 +1710,8 @@ def test_command_job_node_services_in_pipeline_with_no_component(self):
 
         # rest object of node in pipeline should be pure dict
         assert job_rest_obj.properties.jobs["hello_world_component_inline"]["services"] == {
-            "my_jupyter": {
-                "job_service_type": "Jupyter",
+            "my_ssh": {
+                "job_service_type": "SSH",
             },
             "my_tensorboard": {
                 "job_service_type": "TensorBoard",

diff --git a/.../compute/e2etests/test_compute.pyTestComputetest_compute_instance_stop_start_restart.json b/.../compute/e2etests/test_compute.pyTestComputetest_compute_instance_stop_start_restart.json
@@ -695,10 +695,6 @@
               "privateIpAddress": "10.0.0.4"
             },
             "applications": [
-              {
-                "displayName": "Jupyter",
-                "endpointUri": "https://test322425765248.eastus2euap.instances.azureml.ms/tree/"
-              },
               {
                 "displayName": "Jupyter Lab",
                 "endpointUri": "https://test322425765248.eastus2euap.instances.azureml.ms/lab"
@@ -948,10 +944,6 @@
               "privateIpAddress": "10.0.0.4"
             },
             "applications": [
-              {
-                "displayName": "Jupyter",
-                "endpointUri": "https://test322425765248.eastus2euap.instances.azureml.ms/tree/"
-              },
               {
                 "displayName": "Jupyter Lab",
                 "endpointUri": "https://test322425765248.eastus2euap.instances.azureml.ms/lab"
@@ -1237,7 +1229,7 @@
             },
             "applications": [
               {
-                "displayName": "Jupyter",
+                "displayName": "ssh",
                 "endpointUri": "https://test322425765248.eastus2euap.instances.azureml.ms/tree/"
               },
               {
@@ -1315,4 +1307,4 @@
   "Variables": {
     "compute_name": "test322425765248"
   }
-}
+}
diff --git a/...gs/pipeline_job/e2etests/test_pipeline_job.pyTestPipelineJobtest_pipeline_job_create.json b/...gs/pipeline_job/e2etests/test_pipeline_job.pyTestPipelineJobtest_pipeline_job_create.json
@@ -647,15 +647,15 @@
           "experimentName": "my_first_experiment",
           "services": {
             "Tracking": {
-              "jobServiceType": "Tracking",
+              "jobServiceType": "tracking",
               "port": null,
               "endpoint": "azureml://eastus2euap.api.azureml.ms/mlflow/v1.0/subscriptions/00000000-0000-0000-0000-000000000/resourceGroups/00000/providers/Microsoft.MachineLearningServices/workspaces/00000?",
               "status": null,
               "errorMessage": null,
               "properties": null
             },
             "Studio": {
-              "jobServiceType": "Studio",
+              "jobServiceType": "studio",
               "port": null,
               "endpoint": "https://ml.azure.com/runs/test_917642482668?wsid=/subscriptions/00000000-0000-0000-0000-000000000/resourcegroups/00000/workspaces/00000",
               "status": null,
@@ -1174,15 +1174,15 @@
           "experimentName": "my_first_experiment",
           "services": {
             "Tracking": {
-              "jobServiceType": "Tracking",
+              "jobServiceType": "tracking",
               "port": null,
               "endpoint": "azureml://eastus2euap.api.azureml.ms/mlflow/v1.0/subscriptions/00000000-0000-0000-0000-000000000/resourceGroups/00000/providers/Microsoft.MachineLearningServices/workspaces/00000?",
               "status": null,
               "errorMessage": null,
               "properties": null
             },
             "Studio": {
-              "jobServiceType": "Studio",
+              "jobServiceType": "studio",
               "port": null,
               "endpoint": "https://ml.azure.com/runs/test_917642482668?wsid=/subscriptions/00000000-0000-0000-0000-000000000/resourcegroups/00000/workspaces/00000",
               "status": null,
@@ -1294,4 +1294,4 @@
     "new_tag_name": "test_907076112763",
     "new_tag_value": "test_34294107164"
   }
-}
+}
diff --git a/sdk/ml/azure-ai-ml/tests/test_configs/command_job/command_job_inputs_rest.yml b/sdk/ml/azure-ai-ml/tests/test_configs/command_job/command_job_inputs_rest.yml
@@ -7,11 +7,11 @@
       "provisioning_state":"Succeeded",
       "services":{
          "Tracking":{
-            "type":"Tracking",
+            "type":"tracking",
             "endpoint":"azureml://master.api.azureml-test.ms/mlflow/v1.0/subscriptions/d511f82f-71ba-49a4-8233-d7be8a3650f4/resourceGroups/RLTesting/providers/Microsoft.MachineLearningServices/workspaces/AnkitWS?"
          },
          "Studio":{
-            "type":"Studio",
+            "type":"studio",
             "endpoint":"https://ml.azure.com/experiments/mfe-test1/runs/test_773798882091?wsid=/subscriptions/d511f82f-71ba-49a4-8233-d7be8a3650f4/resourcegroups/RLTesting/workspaces/AnkitWS"
          }
       },

diff --git a/sdk/ml/azure-ai-ml/tests/test_configs/command_job/command_job_rest.yml b/sdk/ml/azure-ai-ml/tests/test_configs/command_job/command_job_rest.yml
@@ -7,11 +7,11 @@
       "provisioning_state":"Succeeded",
       "services":{
          "Tracking":{
-            "type":"Tracking",
+            "type":"tracking",
             "endpoint":"azureml://master.api.azureml-test.ms/mlflow/v1.0/subscriptions/d511f82f-71ba-49a4-8233-d7be8a3650f4/resourceGroups/RLTesting/providers/Microsoft.MachineLearningServices/workspaces/AnkitWS?"
          },
          "Studio":{
-            "type":"Studio",
+            "type":"studio",
             "endpoint":"https://ml.azure.com/experiments/mfe-test1/runs/test_617704734544?wsid=/subscriptions/d511f82f-71ba-49a4-8233-d7be8a3650f4/resourcegroups/RLTesting/workspaces/AnkitWS"
          }
       },

diff --git a/sdk/ml/azure-ai-ml/tests/test_configs/command_job/command_job_rest_amltoken_identity.yml b/sdk/ml/azure-ai-ml/tests/test_configs/command_job/command_job_rest_amltoken_identity.yml
@@ -7,11 +7,11 @@
       "provisioning_state":"Succeeded",
       "services":{
          "Tracking":{
-            "type":"Tracking",
+            "type":"tracking",
             "endpoint":"azureml://master.api.azureml-test.ms/mlflow/v1.0/subscriptions/d511f82f-71ba-49a4-8233-d7be8a3650f4/resourceGroups/RLTesting/providers/Microsoft.MachineLearningServices/workspaces/AnkitWS?"
          },
          "Studio":{
-            "type":"Studio",
+            "type":"studio",
             "endpoint":"https://ml.azure.com/experiments/mfe-test1/runs/test_617704734544?wsid=/subscriptions/d511f82f-71ba-49a4-8233-d7be8a3650f4/resourcegroups/RLTesting/workspaces/AnkitWS"
          }
       },

diff --git a/sdk/ml/azure-ai-ml/tests/test_configs/command_job/command_job_rest_msi_identity.yml b/sdk/ml/azure-ai-ml/tests/test_configs/command_job/command_job_rest_msi_identity.yml
@@ -7,11 +7,11 @@
       "provisioning_state":"Succeeded",
       "services":{
          "Tracking":{
-            "type":"Tracking",
+            "type":"tracking",
             "endpoint":"azureml://master.api.azureml-test.ms/mlflow/v1.0/subscriptions/d511f82f-71ba-49a4-8233-d7be8a3650f4/resourceGroups/RLTesting/providers/Microsoft.MachineLearningServices/workspaces/AnkitWS?"
          },
          "Studio":{
-            "type":"Studio",
+            "type":"studio",
             "endpoint":"https://ml.azure.com/experiments/mfe-test1/runs/test_297155932351?wsid=/subscriptions/d511f82f-71ba-49a4-8233-d7be8a3650f4/resourcegroups/RLTesting/workspaces/AnkitWS"
          }
       },

diff --git a/...ure-ai-ml/tests/test_configs/pipeline_jobs/helloworld_pipeline_job_with_node_services.yml b/...ure-ai-ml/tests/test_configs/pipeline_jobs/helloworld_pipeline_job_with_node_services.yml
@@ -18,11 +18,11 @@ jobs:
       environment: azureml:AzureML-sklearn-0.24-ubuntu18.04-py37-cpu:1
       code: "./"
     services:
-      "my_jupyter":
-        job_service_type: "Jupyter" # Jupyter Notebook
+      "my_ssh":
+        job_service_type: "ssh"
       "my_tensorboard":
-        job_service_type: "TensorBoard"
+        job_service_type: "tensor_board"
         properties:
-          logDir: "~/tblog" # where you want to store the TensorBoard output
+          logDir: "~/tblog" # where you want to store the tensor_board output
       "my_jupyterlab":
-        job_service_type: "JupyterLab"
+        job_service_type: "jupyter_lab"