[ml] Fix azure/ai/ml/_internal mypy errors (#33643)

* Add mypy.ini file

* Fix bug 2810296, replace ResourceConfiguration with CodeConfiguration

* Turn on mypy gate

* Update pyproject.toml

* Remove mypy.ini file

* Update exclude section of pyproject.toml

* Exclude _schema folder from mypy check

* Fix _internal errors

* Fix remaining _internal errors

* Fix remaining _internal errors

* Revert _schema fixes, _schema excluded from checks

* Address comment

* Update pyproject.toml to exclude everything but _internal

* Update pyproject.toml to exclude everything but _internal

* Exclude remaining files

* Ignore overload error in inputs_outputs

* Fix typo

* Revert is_deterministic changes

* Fix formatting

* Revert last is_deterministic change

* Fix line too long pylint error

Author: diondrapeck

sdk/ml/azure-ai-ml/azure/ai/ml/_internal/entities/_additional_includes.py

@@ -2,6 +2,7 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # ---------------------------------------------------------
 from pathlib import Path
+from typing import Optional
 
 from ...constants._common import AzureDevopsArtifactsType
 from ...entities._component._additional_includes import AdditionalIncludes
@@ -17,7 +18,7 @@ def includes(self):
         return self.origin_configs
 
     @property
-    def code_path(self) -> Path:
+    def code_path(self) -> Optional[Path]:
         return self.resolved_code_path
 
     @property

sdk/ml/azure-ai-ml/azure/ai/ml/_internal/entities/_input_outputs.py

@@ -91,12 +91,13 @@ def _get_python_builtin_type_str(self) -> str:
         if self._lower_type in [_INPUT_TYPE_ENUM, _INPUT_TYPE_FLOAT]:
             return self._lower_type
         if self._is_primitive_type:
-            return IOConstants.PRIMITIVE_STR_2_TYPE[self._lower_type].__name__
+            return IOConstants.PRIMITIVE_STR_2_TYPE[self._lower_type].__name__  # type: ignore[index]
+            # TODO: Bug 2881900
         return super()._get_python_builtin_type_str()
 
     @overload
     @classmethod
-    def _from_base(cls, _input: None) -> None:
+    def _from_base(cls, _input: None) -> None:  # type: ignore[misc]
         ...
 
     @overload
@@ -125,7 +126,7 @@ def _from_base(cls, _input: Optional[Union[Input, Dict]]) -> Optional["InternalI
             # do force cast directly as there is no new field added in InternalInput
             # need to change the logic if new field is added
             _input.__class__ = InternalInput
-            return _input
+            return _input  # type: ignore[return-value]
         return InternalInput(**_input)
 
 
@@ -180,7 +181,7 @@ def _from_base(cls, _output: Union[Output, Dict]) -> Optional["InternalOutput"]:
             # do force cast directly as there is no new field added in InternalInput
             # need to change the logic if new field is added
             _output.__class__ = InternalOutput
-            return _output
+            return _output  # type: ignore[return-value]
         return InternalOutput(**_output)
 
     def map_pipeline_output_type(self) -> str:

sdk/ml/azure-ai-ml/azure/ai/ml/_internal/entities/command.py

@@ -2,11 +2,11 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # ---------------------------------------------------------
 # pylint: disable=protected-access
-from typing import Dict, List, Union
+from typing import Dict, List, Optional, Union
 
 from marshmallow import INCLUDE, Schema
 
-from ... import MpiDistribution, PyTorchDistribution, TensorFlowDistribution, RayDistribution
+from ... import MpiDistribution, PyTorchDistribution, RayDistribution, TensorFlowDistribution
 from ..._schema import PathAwareSchema
 from ..._schema.core.fields import DistributionField
 from ...entities import CommandJobLimits, JobResourceConfiguration
@@ -34,16 +34,16 @@ def __init__(self, **kwargs):
         self._init = False
 
     @property
-    def compute(self) -> str:
+    def compute(self) -> Optional[str]:
         """Get the compute definition for the command.
 
         :return: The compute definition
-        :rtype: str
+        :rtype: Optional[str]
         """
         return self._compute
 
     @compute.setter
-    def compute(self, value: str):
+    def compute(self, value: str) -> None:
         """Set the compute definition for the command.
 
         :param value: The new compute definition
@@ -52,16 +52,16 @@ def compute(self, value: str):
         self._compute = value
 
     @property
-    def environment(self) -> str:
+    def environment(self) -> Optional[str]:
         """Get the environment definition for the command.
 
         :return: The environment definition
-        :rtype: str
+        :rtype: Optional[str]
         """
         return self._environment
 
     @environment.setter
-    def environment(self, value: str):
+    def environment(self, value: str) -> None:
         """Set the environment definition for the command.
 
         :param value: The new environment definition
@@ -70,16 +70,16 @@ def environment(self, value: str):
         self._environment = value
 
     @property
-    def environment_variables(self) -> Dict[str, str]:
+    def environment_variables(self) -> Optional[Dict[str, str]]:
         """Get the environment variables for the command.
 
         :return: The environment variables
-        :rtype: Dict[str, str]
+        :rtype: Optional[Dict[str, str]]
         """
         return self._environment_variables
 
     @environment_variables.setter
-    def environment_variables(self, value: Dict[str, str]):
+    def environment_variables(self, value: Dict[str, str]) -> None:
         """Set the environment variables for the command.
 
         :param value: The new environment variables

sdk/ml/azure-ai-ml/azure/ai/ml/_internal/entities/component.py

@@ -3,13 +3,14 @@
 # ---------------------------------------------------------
 # pylint: disable=protected-access, redefined-builtin
 # disable redefined-builtin to use id/type as argument name
+import os
 from contextlib import contextmanager
 from os import PathLike
 from pathlib import Path
-from typing import Dict, Iterable, List, Optional, Union
+from typing import Any, Dict, Iterable, List, Optional, Union
 from uuid import UUID
 
-import yaml
+import yaml  # type: ignore[import]
 from marshmallow import Schema
 
 from ... import Input, Output
@@ -20,7 +21,7 @@
 from ...constants._common import DefaultOpenEncoding
 from ...entities import Component
 from ...entities._assets import Code
-from ...entities._component._additional_includes import AdditionalIncludesMixin, AdditionalIncludes
+from ...entities._component._additional_includes import AdditionalIncludes, AdditionalIncludesMixin
 from ...entities._component.code import ComponentIgnoreFile
 from ...entities._job.distribution import DistributionConfiguration
 from ...entities._system_data import SystemData
@@ -86,7 +87,7 @@ def __init__(
         successful_return_code: Optional[str] = None,
         inputs: Optional[Dict] = None,
         outputs: Optional[Dict] = None,
-        code: Optional[str] = None,
+        code: Optional[Union[str, os.PathLike]] = None,
         environment: Optional[Dict] = None,
         environment_variables: Optional[Dict] = None,
         command: Optional[str] = None,
@@ -200,11 +201,13 @@ def _get_additional_includes_field_name(cls) -> str:
 
     def _get_all_additional_includes_configs(self) -> List:
         # internal components must have a source path
-        return self._read_additional_include_configs(Path(self._source_path))
+        return self._read_additional_include_configs(Path(self._source_path))  # type: ignore[arg-type]
+        # TODO: Bug 2881943
 
     def _get_base_path_for_code(self) -> Path:
         # internal components must have a source path
-        return Path(self._source_path).parent
+        return Path(self._source_path).parent  # type: ignore[arg-type]
+        # TODO: Bug 2881943
 
     def _get_origin_code_value(self) -> Union[str, PathLike, None]:
         return super()._get_origin_code_value() or "."
@@ -276,8 +279,9 @@ def _from_rest_object_to_init_params(cls, obj: ComponentVersion) -> Dict:
         return init_kwargs
 
     def _to_rest_object(self) -> ComponentVersion:
-        component = convert_ordered_dict_to_dict(self._to_dict())
-        component["_source"] = self._source
+        component: Union[Dict[Any, Any], List[Any]] = convert_ordered_dict_to_dict(self._to_dict())
+        component["_source"] = self._source  # type: ignore[call-overload]
+        # TODO: 2883063
 
         properties = ComponentVersionProperties(
             component_spec=component,
@@ -310,7 +314,7 @@ def _get_snapshot_id(
         snapshot_id = str(UUID(curr_root.hexdigest_hash[::4]))
         return snapshot_id
 
-    @contextmanager
+    @contextmanager  # type: ignore[arg-type]
     def _try_build_local_code(self) -> Iterable[Code]:
         """Build final code when origin code is a local code.
         Will merge code path with additional includes into a temp folder if additional includes is specified.
@@ -319,6 +323,8 @@ def _try_build_local_code(self) -> Iterable[Code]:
         :return: The code instance
         :rtype: Iterable[Code]
         """
+
+        tmp_code_dir: Path
         # origin code value of internal component will never be None. check _get_origin_code_value for details
         with self._generate_additional_includes_obj().merge_local_code_and_additional_includes() as tmp_code_dir:
             # use absolute path in case temp folder & work dir are in different drive
@@ -334,8 +340,12 @@ def _try_build_local_code(self) -> Iterable[Code]:
             # additional includes config file itself should be ignored
             rebased_ignore_file = ComponentIgnoreFile(
                 tmp_code_dir,
-                additional_includes_file_name=Path(self._source_path).with_suffix(_ADDITIONAL_INCLUDES_SUFFIX).name,
+                additional_includes_file_name=Path(self._source_path)
+                .with_suffix(_ADDITIONAL_INCLUDES_SUFFIX)
+                .name  # type: ignore[arg-type]
+                # TODO: Bug 2881943
             )
+
             # Use the snapshot id in ml-components as code name to enable anonymous
             # component reuse from ml-component runs.
             # calculate snapshot id here instead of inside InternalCode to ensure that

sdk/ml/azure-ai-ml/azure/ai/ml/_internal/entities/node.py

@@ -65,7 +65,8 @@ def __init__(
         BaseNode.__init__(
             self,
             type=type,
-            component=component,
+            component=component,  # type: ignore[arg-type]
+            # TODO: Bug 2881892
             inputs=inputs,
             outputs=outputs,
             compute=compute,

sdk/ml/azure-ai-ml/azure/ai/ml/_internal/entities/runsettings/ai_super_computer_configuration.py

@@ -2,13 +2,13 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # ---------------------------------------------------------
 
-from typing import Dict, List, Optional
+from typing import Any, Dict, List, Optional
 
 from ....entities._job.job_resource_configuration import BaseProperty
 
 
 class PascalCaseProperty(BaseProperty):
-    _KEY_MAPPING = {}
+    _KEY_MAPPING: Dict[str, Any] = {}
 
     def items(self):
         result = []

sdk/ml/azure-ai-ml/azure/ai/ml/_internal/entities/spark.py

@@ -40,7 +40,7 @@ def __init__(
         dynamic_allocation_min_executors: Optional[int] = None,
         dynamic_allocation_max_executors: Optional[int] = None,
         conf: Optional[Dict[str, str]] = None,
-        args: Optional[Dict] = None,
+        args: Optional[str] = None,
         **kwargs,
     ):
         SparkJobEntryMixin.__init__(self, entry=entry, **kwargs)
@@ -90,12 +90,12 @@ def __init__(
     def _create_schema_for_validation(cls, context) -> Union[PathAwareSchema, Schema]:
         return InternalSparkComponentSchema(context=context)
 
-    @property
-    def environment(self) -> Union[str, Environment, dict]:
+    @property  # type: ignore[override]
+    def environment(self) -> Optional[Union[Environment, str]]:
         """Get the environment of the component.
 
         :return: The environment of the component.
-        :rtype: Union[str, Environment, dict]
+        :rtype: Optional[Union[Environment, str]]]
         """
         if isinstance(self._environment, Environment) and self._environment.image is None:
             return Environment(conda_file=self._environment.conda_file, image=DUMMY_IMAGE)
@@ -131,11 +131,11 @@ def environment(self, value):
             raise ValueError(f"Unsupported environment type: {type(value)}")
 
     @property
-    def jars(self) -> List[str]:
+    def jars(self) -> Optional[List[str]]:
         """Get the jars of the component.
 
         :return: The jars of the component.
-        :rtype: List[str]
+        :rtype: Optional[List[str]]
         """
         return self._jars
 
@@ -153,11 +153,11 @@ def jars(self, value: Union[str, List[str]]):
         self._jars = value
 
     @property
-    def py_files(self) -> List[str]:
+    def py_files(self) -> Optional[List[str]]:
         """Get the py_files of the component.
 
         :return: The py_files of the component.
-        :rtype: List[str]
+        :rtype: Optional[List[str]]
         """
         return self._py_files
 

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/_datastore/azure_storage.py

@@ -47,7 +47,7 @@ class AzureFileSchema(AzureStorageSchema):
     )
 
     @post_load
-    def make(self, data: Dict[str, Any], **kwargs) -> "AzureFileDatastore":
+    def make(self, data: Dict[str, Any], **kwargs) -> "AzureFileDatastore":  # type: ignore[name-defined]
         from azure.ai.ml.entities import AzureFileDatastore
 
         return AzureFileDatastore(**data)
@@ -69,7 +69,7 @@ class AzureBlobSchema(AzureStorageSchema):
     )
 
     @post_load
-    def make(self, data: Dict[str, Any], **kwargs) -> "AzureBlobDatastore":
+    def make(self, data: Dict[str, Any], **kwargs) -> "AzureBlobDatastore":  # type: ignore[name-defined]
         from azure.ai.ml.entities import AzureBlobDatastore
 
         return AzureBlobDatastore(**data)

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/_datastore/one_lake.py

@@ -6,10 +6,9 @@
 
 from typing import Any, Dict
 
-from marshmallow import fields, post_load, Schema
+from marshmallow import Schema, fields, post_load
 
 from azure.ai.ml._restclient.v2023_04_01_preview.models import DatastoreType, OneLakeArtifactType
-
 from azure.ai.ml._schema.core.fields import NestedField, PathAwareSchema, StringTransformedEnum, UnionField
 from azure.ai.ml._utils.utils import camel_to_snake
 

sdk/ml/azure-ai-ml/azure/ai/ml/_utils/_arm_id_utils.py

@@ -330,8 +330,8 @@ def parse_name_version(name: str) -> Tuple[str, Optional[str]]:
     return name, ":".join(version)
 
 
-def parse_name_label(name: str) -> Tuple[str, Optional[str]]:
-    if name.find("/") != -1 and name[0] != "/":
+def parse_name_label(name: Optional[str]) -> Tuple[str, Optional[str]]:
+    if not name or (name.find("/") != -1 and name[0] != "/"):
         raise ValidationException(
             f"Could not parse {name}. If providing an ARM id, it should start with a '/'.",
             no_personal_data_message=f"Could not parse {name}.",