[ML] Restrict JobServiceType to jupyter_lab, ssh, tensor_board, vs_code and perform required conversion before sending to rest API

sdk/ml/azure-ai-ml/azure/ai/ml/constants/_job/job.py

@@ -72,3 +72,26 @@ 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:
+        JUPYTER_LAB = "jupyter_lab"
+        SSH = "ssh"
+        TENSOR_BOARD = "tensor_board"
+        VS_CODE = "vs_code"
+
+    class RestNames:
+        JUPYTER_LAB = "JupyterLab"
+        SSH = "SSH"
+        TENSOR_BOARD = "TensorBoard"
+        VS_CODE = "VSCode"
+
+    ENTITY_TO_REST = {
+        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()}

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,9 @@ def __init__(
         self,
         *,
         endpoint: Optional[str] = None,
-        job_service_type: Optional[Literal["JupyterLab", "SSH", "TensorBoard", "VSCode"]] = None,
+        # TODO: Is VS Code allowed here?
+        # TODO: As this is a breaking change, should we still allow the old format (any-case) and deprecate in the next release?
+        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 +54,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 +79,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"{list(JobServiceTypeNames.ENTITY_TO_REST.keys())}, 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 +105,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,

sdk/ml/azure-ai-ml/tests/dsl/unittests/test_command_builder.py

@@ -839,9 +839,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",
                 },
@@ -874,14 +874,14 @@ 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": "TensorBoard",
+                "job_service_type": "tensor_board",
                 "properties": {
                     "logDir": "~/tblog",
                 },
             },
-            "my_jupyterlab": {"job_service_type": "JupyterLab"},
+            "my_jupyterlab": {"job_service_type": "jupyter_lab"},
         }
         node = command(
             name="interactive-command-job",
@@ -903,7 +903,7 @@ def test_command_services(self) -> None:
         assert node_rest_obj["services"] == 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",

sdk/ml/azure-ai-ml/tests/dsl/unittests/test_dsl_pipeline.py

@@ -3700,14 +3700,14 @@ def root_pipeline(component_in_number: int, component_in_path: 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": "TensorBoard",
+                "job_service_type": "tensor_board",
                 "properties": {
                     "logDir": "~/tblog",
                 },
             },
-            "my_jupyterlab": {"job_service_type": "JupyterLab"},
+            "my_jupyterlab": {"job_service_type": "jupyter_lab"},
         }
 
         command_func = command(
@@ -3749,7 +3749,7 @@ def sample_pipeline():
             assert isinstance(service, JobService)
 
         # test set services in pipeline
-        new_services = {"my_jupyter": {"job_service_type": "Jupyter"}}
+        new_services = {"my_ssh": {"job_service_type": "ssh"}}
 
         @dsl.pipeline()
         def sample_pipeline_with_new_services():

sdk/ml/azure-ai-ml/tests/pipeline_job/unittests/test_pipeline_job_schema.py

@@ -1686,15 +1686,15 @@ 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",
+                "job_service_type": "tensor_board",
                 "properties": {"logDir": "~/tblog"},
             },
             "my_jupyterlab": {
-                "job_service_type": "JupyterLab",
+                "job_service_type": "jupyter_lab",
             },
         }
 
@@ -1710,15 +1710,15 @@ 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",
+                "job_service_type": "tensor_board",
                 "properties": {"logDir": "~/tblog"},
             },
             "my_jupyterlab": {
-                "job_service_type": "JupyterLab",
+                "job_service_type": "jupyter_lab",
             },
         }
 

sdk/ml/azure-ai-ml/tests/recordings/compute/e2etests/test_compute.pyTestComputetest_compute_instance_stop_start_restart.json

@@ -696,7 +696,7 @@
             },
             "applications": [
               {
-                "displayName": "Jupyter",
+                "displayName": "ssh",
                 "endpointUri": "https://test322425765248.eastus2euap.instances.azureml.ms/tree/"
               },
               {
@@ -949,7 +949,7 @@
             },
             "applications": [
               {
-                "displayName": "Jupyter",
+                "displayName": "ssh",
                 "endpointUri": "https://test322425765248.eastus2euap.instances.azureml.ms/tree/"
               },
               {
@@ -1237,7 +1237,7 @@
             },
             "applications": [
               {
-                "displayName": "Jupyter",
+                "displayName": "ssh",
                 "endpointUri": "https://test322425765248.eastus2euap.instances.azureml.ms/tree/"
               },
               {
@@ -1315,4 +1315,4 @@
   "Variables": {
     "compute_name": "test322425765248"
   }
-}
+}

sdk/ml/azure-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: "Jupyter"
       "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"

sdk/ml/azure-ai-ml/tests/test_configs/pipeline_jobs/helloworld_pipeline_job_with_node_services_inline_job.yml

@@ -13,11 +13,11 @@ jobs:
     command: echo Hello World & sleep 1h
     environment: azureml:AzureML-sklearn-0.24-ubuntu18.04-py37-cpu:1
     services:
-      "my_jupyter":
-        job_service_type: "Jupyter" # Jupyter Notebook
+      "my_ssh":
+        job_service_type: "Jupyter"
       "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"