From bc07d34355c410f1e16842f7bcdc3235b91ac1d3 Mon Sep 17 00:00:00 2001 From: Zubairuddin Mohammed Date: Fri, 21 Nov 2025 19:34:24 +0530 Subject: [PATCH 1/2] next-mypy fixes for automl in azure-ai-ml sdk --- .../image/automl_image_classification_base.py | 4 ++-- .../automl_image_object_detection_base.py | 2 +- .../_job/automl/nlp/automl_nlp_job.py | 2 +- .../_job/automl/tabular/classification_job.py | 2 +- .../automl/tabular/featurization_settings.py | 4 ++-- .../_job/automl/tabular/forecasting_job.py | 4 ++-- .../_job/automl/tabular/regression_job.py | 4 ++-- .../entities/_job/automl/training_settings.py | 24 ++++++++++++++----- 8 files changed, 29 insertions(+), 17 deletions(-) diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/image/automl_image_classification_base.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/image/automl_image_classification_base.py index 218dc91cb395..ddc251815390 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/image/automl_image_classification_base.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/image/automl_image_classification_base.py @@ -60,11 +60,11 @@ def training_parameters(self) -> Optional[ImageModelSettingsClassification]: return self._training_parameters @training_parameters.setter - def training_parameters(self, value: Union[Dict, ImageModelSettingsClassification]) -> None: + def training_parameters(self, value: Optional[Union[Dict, ImageModelSettingsClassification]]) -> None: """Setting Image training parameters for AutoML Image Classification and Image Classification Multilabel tasks. :param value: Training parameters for AutoML Image Classification and Image Classification Multilabel tasks. - :type value: Union[Dict, ~azure.ai.ml.automl.ImageModelSettingsClassification] + :type value: Optional[Union[Dict, ~azure.ai.ml.automl.ImageModelSettingsClassification]] :raises ~azure.ml.exceptions.ValidationException if value is not a dictionary or ImageModelSettingsClassification. :return: None diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/image/automl_image_object_detection_base.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/image/automl_image_object_detection_base.py index db0c7bc6c660..b7c39a2453c7 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/image/automl_image_object_detection_base.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/image/automl_image_object_detection_base.py @@ -64,7 +64,7 @@ def training_parameters(self) -> Optional[ImageModelSettingsObjectDetection]: return self._training_parameters @training_parameters.setter - def training_parameters(self, value: Union[Dict, ImageModelSettingsObjectDetection]) -> None: + def training_parameters(self, value: Optional[Union[Dict, ImageModelSettingsObjectDetection]]) -> None: if value is None: self._training_parameters = None elif isinstance(value, ImageModelSettingsObjectDetection): diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/nlp/automl_nlp_job.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/nlp/automl_nlp_job.py index f0b3baa852aa..8a6f7575b9d3 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/nlp/automl_nlp_job.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/nlp/automl_nlp_job.py @@ -100,7 +100,7 @@ def training_parameters(self) -> Optional[NlpFixedParameters]: return self._training_parameters @training_parameters.setter - def training_parameters(self, value: Union[Dict, NlpFixedParameters]) -> None: + def training_parameters(self, value: Optional[Union[Dict, NlpFixedParameters]]) -> None: if value is None: self._training_parameters = None elif isinstance(value, NlpFixedParameters): diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/tabular/classification_job.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/tabular/classification_job.py index 6f5ab2719cb3..daaf03538f87 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/tabular/classification_job.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/tabular/classification_job.py @@ -124,7 +124,7 @@ def training(self) -> ClassificationTrainingSettings: return self._training or ClassificationTrainingSettings() @training.setter - def training(self, value: Union[Dict, ClassificationTrainingSettings]) -> None: # pylint: disable=unused-argument + def training(self, value: Union[Dict, TrainingSettings]) -> None: # pylint: disable=unused-argument ... def _to_rest_object(self) -> JobBase: diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/tabular/featurization_settings.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/tabular/featurization_settings.py index 6ef2332ee370..ff861faddee6 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/tabular/featurization_settings.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/tabular/featurization_settings.py @@ -101,7 +101,7 @@ def transformer_params(self) -> Optional[Dict[str, List[ColumnTransformer]]]: return self._transformer_params @transformer_params.setter - def transformer_params(self, value: Dict[str, List[ColumnTransformer]]) -> None: + def transformer_params(self, value: Optional[Dict[str, List[ColumnTransformer]]]) -> None: self._transformer_params = ( None if not value @@ -114,7 +114,7 @@ def blocked_transformers(self) -> Optional[List[Union[BlockedTransformers, str]] return self._blocked_transformers @blocked_transformers.setter - def blocked_transformers(self, blocked_transformers_list: List[Union[BlockedTransformers, str]]) -> None: + def blocked_transformers(self, blocked_transformers_list: Optional[List[Union[BlockedTransformers, str]]]) -> None: self._blocked_transformers = ( None if blocked_transformers_list is None diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/tabular/forecasting_job.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/tabular/forecasting_job.py index 9bd10b19875d..bec515dbafec 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/tabular/forecasting_job.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/tabular/forecasting_job.py @@ -20,7 +20,7 @@ from azure.ai.ml.entities._job.automl.tabular.featurization_settings import TabularFeaturizationSettings from azure.ai.ml.entities._job.automl.tabular.forecasting_settings import ForecastingSettings from azure.ai.ml.entities._job.automl.tabular.limit_settings import TabularLimitSettings -from azure.ai.ml.entities._job.automl.training_settings import ForecastingTrainingSettings +from azure.ai.ml.entities._job.automl.training_settings import ForecastingTrainingSettings, TrainingSettings from azure.ai.ml.entities._util import load_from_dict @@ -101,7 +101,7 @@ def training(self) -> ForecastingTrainingSettings: return self._training or ForecastingTrainingSettings() @training.setter - def training(self, value: Union[Dict, ForecastingTrainingSettings]) -> None: # pylint: disable=unused-argument + def training(self, value: Union[Dict, TrainingSettings]) -> None: # pylint: disable=unused-argument ... @property diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/tabular/regression_job.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/tabular/regression_job.py index 3531e52c16f3..ba0d215b52be 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/tabular/regression_job.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/tabular/regression_job.py @@ -16,7 +16,7 @@ from azure.ai.ml.entities._credentials import _BaseJobIdentityConfiguration from azure.ai.ml.entities._job._input_output_helpers import from_rest_data_outputs, to_rest_data_outputs from azure.ai.ml.entities._job.automl.tabular import AutoMLTabular, TabularFeaturizationSettings, TabularLimitSettings -from azure.ai.ml.entities._job.automl.training_settings import RegressionTrainingSettings +from azure.ai.ml.entities._job.automl.training_settings import RegressionTrainingSettings, TrainingSettings from azure.ai.ml.entities._util import load_from_dict @@ -74,7 +74,7 @@ def training(self) -> RegressionTrainingSettings: return self._training or RegressionTrainingSettings() @training.setter - def training(self, value: Union[Dict, RegressionTrainingSettings]) -> None: # pylint: disable=unused-argument + def training(self, value: Union[Dict, TrainingSettings]) -> None: # pylint: disable=unused-argument ... def _to_rest_object(self) -> JobBase: diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/training_settings.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/training_settings.py index 97bc7e171963..ef0eee38db28 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/training_settings.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/automl/training_settings.py @@ -190,7 +190,9 @@ def allowed_training_algorithms(self) -> Optional[List]: return self._allowed_training_algorithms @allowed_training_algorithms.setter - def allowed_training_algorithms(self, allowed_model_list: Union[List[str], List[ClassificationModels]]) -> None: + def allowed_training_algorithms( + self, allowed_model_list: Optional[Union[List[str], List[ClassificationModels]]] + ) -> None: self._allowed_training_algorithms = ( None if allowed_model_list is None @@ -202,7 +204,9 @@ def blocked_training_algorithms(self) -> Optional[List]: return self._blocked_training_algorithms @blocked_training_algorithms.setter - def blocked_training_algorithms(self, blocked_model_list: Union[List[str], List[ClassificationModels]]) -> None: + def blocked_training_algorithms( + self, blocked_model_list: Optional[Union[List[str], List[ClassificationModels]]] + ) -> None: self._blocked_training_algorithms = ( None if blocked_model_list is None @@ -253,7 +257,9 @@ def allowed_training_algorithms(self) -> Optional[List]: return self._allowed_training_algorithms @allowed_training_algorithms.setter - def allowed_training_algorithms(self, allowed_model_list: Union[List[str], List[ForecastingModels]]) -> None: + def allowed_training_algorithms( + self, allowed_model_list: Optional[Union[List[str], List[ForecastingModels]]] + ) -> None: self._allowed_training_algorithms = ( None if allowed_model_list is None else [ForecastingModels[camel_to_snake(o)] for o in allowed_model_list] ) @@ -263,7 +269,9 @@ def blocked_training_algorithms(self) -> Optional[List]: return self._blocked_training_algorithms @blocked_training_algorithms.setter - def blocked_training_algorithms(self, blocked_model_list: Union[List[str], List[ForecastingModels]]) -> None: + def blocked_training_algorithms( + self, blocked_model_list: Optional[Union[List[str], List[ForecastingModels]]] + ) -> None: self._blocked_training_algorithms = ( None if blocked_model_list is None else [ForecastingModels[camel_to_snake(o)] for o in blocked_model_list] ) @@ -312,7 +320,9 @@ def allowed_training_algorithms(self) -> Optional[List]: return self._allowed_training_algorithms @allowed_training_algorithms.setter - def allowed_training_algorithms(self, allowed_model_list: Union[List[str], List[ForecastingModels]]) -> None: + def allowed_training_algorithms( + self, allowed_model_list: Optional[Union[List[str], List[ForecastingModels]]] + ) -> None: self._allowed_training_algorithms = ( None if allowed_model_list is None else [RegressionModels[camel_to_snake(o)] for o in allowed_model_list] ) @@ -322,7 +332,9 @@ def blocked_training_algorithms(self) -> Optional[List]: return self._blocked_training_algorithms @blocked_training_algorithms.setter - def blocked_training_algorithms(self, blocked_model_list: Union[List[str], List[ForecastingModels]]) -> None: + def blocked_training_algorithms( + self, blocked_model_list: Optional[Union[List[str], List[ForecastingModels]]] + ) -> None: self._blocked_training_algorithms = ( None if blocked_model_list is None else [RegressionModels[camel_to_snake(o)] for o in blocked_model_list] ) From d1cc8e3a4a158ab3222b1b0e88ec1fe6a4597b7c Mon Sep 17 00:00:00 2001 From: Zubairuddin Mohammed Date: Sat, 22 Nov 2025 19:28:53 +0530 Subject: [PATCH 2/2] remainging next mypy issues fix --- .../azure/ai/ml/_internal/entities/spark.py | 4 ++-- .../ai/ml/entities/_assets/_artifacts/data.py | 2 +- .../azure/ai/ml/entities/_assets/asset.py | 4 ++-- .../ai/ml/entities/_assets/environment.py | 4 ++-- .../ai/ml/entities/_builders/base_node.py | 4 ++-- .../azure/ai/ml/entities/_builders/command.py | 23 ++++++++++++------ .../entities/_builders/control_flow_node.py | 2 +- .../ai/ml/entities/_builders/data_transfer.py | 2 +- .../ai/ml/entities/_builders/parallel.py | 21 +++++++++------- .../ai/ml/entities/_builders/pipeline.py | 4 ++-- .../azure/ai/ml/entities/_builders/spark.py | 8 ++++--- .../azure/ai/ml/entities/_builders/sweep.py | 20 +++++++++------- .../entities/_component/command_component.py | 4 ++-- .../ai/ml/entities/_component/component.py | 4 ++-- .../entities/_component/parallel_component.py | 8 +++---- .../ml/entities/_indexes/data_index_func.py | 9 ++++--- .../entities/_inputs_outputs/external_data.py | 2 +- .../azure/ai/ml/entities/_job/job_limits.py | 4 ++-- .../_job/job_resource_configuration.py | 6 ++--- .../ml/entities/_job/parameterized_command.py | 14 ++++++++++- .../ai/ml/entities/_job/pipeline/_io/base.py | 5 ++-- .../ml/entities/_job/pipeline/pipeline_job.py | 4 ++-- .../ml/entities/_job/spark_job_entry_mixin.py | 2 +- .../azure-ai-ml/azure/ai/ml/entities/_util.py | 8 +++---- .../ai/ml/operations/_data_operations.py | 2 +- .../ml/operations/_environment_operations.py | 2 +- .../operations/_feature_store_operations.py | 13 ++++++---- .../azure/ai/ml/operations/_job_operations.py | 24 ++++++++++--------- .../ai/ml/operations/_local_job_invoker.py | 8 +++---- .../ai/ml/operations/_model_operations.py | 3 +-- 30 files changed, 128 insertions(+), 92 deletions(-) diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/_internal/entities/spark.py b/sdk/ml/azure-ai-ml/azure/ai/ml/_internal/entities/spark.py index 345fa5f24cdd..7acba815a2e7 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/_internal/entities/spark.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/_internal/entities/spark.py @@ -140,11 +140,11 @@ def jars(self) -> Optional[List[str]]: return self._jars @jars.setter - def jars(self, value: Union[str, List[str]]): + def jars(self, value: Optional[Union[str, List[str]]]): """Set the jars of the component. :param value: The jars of the component. - :type value: Union[str, List[str]] + :type value: Optional[Union[str, List[str]]] :return: No return :rtype: None """ diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_assets/_artifacts/data.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_assets/_artifacts/data.py index 710e959a4c26..4575b67e6aef 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_assets/_artifacts/data.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_assets/_artifacts/data.py @@ -120,7 +120,7 @@ def path(self) -> Optional[Union[Path, str, PathLike]]: return self._path @path.setter - def path(self, value: str) -> None: + def path(self, value: Optional[Union[str, PathLike]]) -> None: # Call the parent setter to resolve the path with base_path if it was a local path # TODO: Bug Item number: 2883424 super(Data, type(self)).path.fset(self, value) # type: ignore diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_assets/asset.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_assets/asset.py index 3ef0824138a7..cd09bcb509e3 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_assets/asset.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_assets/asset.py @@ -85,11 +85,11 @@ def version(self) -> Optional[str]: return self._version @version.setter - def version(self, value: str) -> None: + def version(self, value: Optional[str]) -> None: """Sets the asset version. :param value: The asset version. - :type value: str + :type value: Optional[str] :raises ValidationException: Raised if value is not a string. """ if value: diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_assets/environment.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_assets/environment.py index dca35fcf5c01..67f1756158bd 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_assets/environment.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_assets/environment.py @@ -284,7 +284,7 @@ def _from_container_rest_object(cls, env_container_rest_object: EnvironmentConta # Setting version to None since if version is not provided it is defaulted to "1". # This should go away once container concept is finalized. - env.version = None + env.version = None # type: ignore[assignment] return env def _to_arm_resource_param(self, **kwargs: Any) -> Dict: # pylint: disable=unused-argument @@ -404,7 +404,7 @@ def _localize(self, base_path: str) -> None: self._creation_context = None self._base_path = base_path if self._is_anonymous: - self.name, self.version = None, None + self.name, self.version = None, None # type: ignore[assignment] # TODO: Remove _DockerBuild and _DockerConfiguration classes once local endpoint moves to using updated env diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/base_node.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/base_node.py index 98eba6a5369e..b1379345dd7c 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/base_node.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/base_node.py @@ -205,11 +205,11 @@ def name(self) -> Optional[str]: return self._name @name.setter - def name(self, value: str) -> None: + def name(self, value: Optional[str]) -> None: """Set the name of the node. :param value: The name to set for the node. - :type value: str + :type value: Optional[str] :return: None """ # when name is not lower case, lower it to make sure it's a valid node name diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/command.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/command.py index 0073307c2b99..317e0b48777f 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/command.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/command.py @@ -260,13 +260,22 @@ def distribution( @distribution.setter def distribution( self, - value: Union[Dict, PyTorchDistribution, TensorFlowDistribution, MpiDistribution, RayDistribution], + value: Optional[ + Union[ + Dict, + PyTorchDistribution, + TensorFlowDistribution, + MpiDistribution, + RayDistribution, + DistributionConfiguration, + ] + ], ) -> None: """Sets the configuration for the distributed command component or job. :param value: The configuration for distributed jobs. - :type value: Union[dict, ~azure.ai.ml.PyTorchDistribution, ~azure.ai.ml.MpiDistribution, - ~azure.ai.ml.TensorFlowDistribution, ~azure.ai.ml.RayDistribution] + :type value: Optional[Union[dict, ~azure.ai.ml.PyTorchDistribution, ~azure.ai.ml.MpiDistribution, + ~azure.ai.ml.TensorFlowDistribution, ~azure.ai.ml.RayDistribution]] """ if isinstance(value, dict): dist_schema = UnionField( @@ -309,11 +318,11 @@ def queue_settings(self) -> Optional[QueueSettings]: return self._queue_settings @queue_settings.setter - def queue_settings(self, value: Union[Dict, QueueSettings]) -> None: + def queue_settings(self, value: Optional[Union[Dict, QueueSettings]]) -> None: """Sets the queue settings for the command component or job. :param value: The queue settings for the command component or job. - :type value: Union[dict, ~azure.ai.ml.entities.QueueSettings] + :type value: Optional[Union[dict, ~azure.ai.ml.entities.QueueSettings]] """ if isinstance(value, dict): value = QueueSettings(**value) @@ -373,7 +382,7 @@ def services( @services.setter def services( self, - value: Dict, + value: Optional[Dict[str, Any]], ) -> None: """Sets the interactive services for the node. @@ -951,7 +960,7 @@ def __call__(self, *args: Any, **kwargs: Any) -> "Command": node.distribution = copy.deepcopy(self.distribution) node.resources = copy.deepcopy(self.resources) node.queue_settings = copy.deepcopy(self.queue_settings) - node.services = copy.deepcopy(self.services) + node.services = copy.deepcopy(self.services) if self.services is not None else None node.identity = copy.deepcopy(self.identity) return node msg = "Command can be called as a function only when referenced component is {}, currently got {}." diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/control_flow_node.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/control_flow_node.py index c757a1e4ba74..a55d9cbef8e5 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/control_flow_node.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/control_flow_node.py @@ -165,6 +165,6 @@ def _is_loop_node_dict(obj: Any) -> bool: def _from_rest_object(cls, obj: dict, pipeline_jobs: dict) -> "LoopNode": from azure.ai.ml.entities._job.pipeline._load_component import pipeline_node_factory - node_type = obj.get(CommonYamlFields.TYPE, None) + node_type = cast(str, obj.get(CommonYamlFields.TYPE, None)) load_from_rest_obj_func = pipeline_node_factory.get_load_from_rest_object_func(_type=node_type) return load_from_rest_obj_func(obj, pipeline_jobs) # type: ignore diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/data_transfer.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/data_transfer.py index 83e88a481083..62833dd4a97a 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/data_transfer.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/data_transfer.py @@ -493,7 +493,7 @@ def sink(self) -> Optional[Union[Dict, Database, FileSystem]]: return self._sink @sink.setter - def sink(self, value: Union[Dict, Database, FileSystem]) -> None: + def sink(self, value: Optional[Union[Dict, Database, FileSystem]]) -> None: self._sink = _build_source_sink(value) @classmethod diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/parallel.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/parallel.py index db1de79730c0..e35f0fd41eac 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/parallel.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/parallel.py @@ -252,11 +252,12 @@ def resources(self) -> Optional[JobResourceConfiguration]: return self._resources @resources.setter - def resources(self, value: Union[JobResourceConfiguration, Dict]) -> None: + def resources(self, value: Optional[Union[JobResourceConfiguration, Dict]]) -> None: """Set the resource configuration for the parallel job. :param value: The resource configuration for the parallel job. - :type value: ~azure.ai.ml.entities._job.job_resource_configuration.JobResourceConfiguration or dict + :type value: Optional[Union[ + ~azure.ai.ml.entities._job.job_resource_configuration.JobResourceConfiguration, dict]] """ if isinstance(value, dict): value = JobResourceConfiguration(**value) @@ -316,11 +317,11 @@ def task(self) -> Optional[ParallelTask]: return self._task # type: ignore @task.setter - def task(self, value: Union[ParallelTask, Dict]) -> None: + def task(self, value: Optional[Union[ParallelTask, Dict]]) -> None: """Set the parallel task. :param value: The parallel task. - :type value: ~azure.ai.ml.entities._job.parallel.parallel_task.ParallelTask or dict + :type value: Optional[Union[~azure.ai.ml.entities._job.parallel.parallel_task.ParallelTask, dict]] """ # base path should be reset if task is set via sdk self._base_path: Optional[Union[str, os.PathLike]] = None @@ -360,16 +361,18 @@ def set_resources( if self.resources is None: self.resources = JobResourceConfiguration() + # Use local variable for type narrowing + resources = self.resources if instance_type is not None: - self.resources.instance_type = instance_type + resources.instance_type = instance_type if instance_count is not None: - self.resources.instance_count = instance_count + resources.instance_count = instance_count if properties is not None: - self.resources.properties = properties + resources.properties = properties if docker_args is not None: - self.resources.docker_args = docker_args + resources.docker_args = docker_args if shm_size is not None: - self.resources.shm_size = shm_size + resources.shm_size = shm_size # Save the resources to internal component as well, otherwise calling sweep() will loose the settings if isinstance(self.component, Component): diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/pipeline.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/pipeline.py index 188d90447f40..4925fb0327d5 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/pipeline.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/pipeline.py @@ -103,11 +103,11 @@ def settings(self) -> Optional[PipelineJobSettings]: return self._settings @settings.setter - def settings(self, value: Union[PipelineJobSettings, Dict]) -> None: + def settings(self, value: Optional[Union[PipelineJobSettings, Dict]]) -> None: """Set the settings of the pipeline. :param value: The settings of the pipeline. - :type value: Union[~azure.ai.ml.entities.PipelineJobSettings, dict] + :type value: Optional[Union[~azure.ai.ml.entities.PipelineJobSettings, dict]] :raises TypeError: If the value is not an instance of PipelineJobSettings or a dict. """ if value is not None: diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/spark.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/spark.py index e72f1334a937..08b759e69059 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/spark.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/spark.py @@ -292,13 +292,15 @@ def identity( @identity.setter def identity( self, - value: Union[Dict[str, str], ManagedIdentityConfiguration, AmlTokenConfiguration, UserIdentityConfiguration], + value: Optional[ + Union[Dict[str, str], ManagedIdentityConfiguration, AmlTokenConfiguration, UserIdentityConfiguration] + ], ) -> None: """Sets the identity that the Spark job will use while running on compute. :param value: The identity that the Spark job will use while running on compute. - :type value: Union[Dict[str, str], ~azure.ai.ml.entities.ManagedIdentityConfiguration, - ~azure.ai.ml.entities.AmlTokenConfiguration, ~azure.ai.ml.entities.UserIdentityConfiguration] + :type value: Optional[Union[Dict[str, str], ~azure.ai.ml.entities.ManagedIdentityConfiguration, + ~azure.ai.ml.entities.AmlTokenConfiguration, ~azure.ai.ml.entities.UserIdentityConfiguration]] """ if isinstance(value, dict): identify_schema = UnionField( diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/sweep.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/sweep.py index 603babbe215c..aefca99b95c7 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/sweep.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/sweep.py @@ -158,7 +158,8 @@ def __init__( **kwargs: Any, ) -> None: # TODO: get rid of self._job_inputs, self._job_outputs once we have general Input - self._job_inputs, self._job_outputs = inputs, outputs + self._job_inputs: Dict[str, Any] = inputs or {} + self._job_outputs: Dict[Any, Any] = outputs or {} kwargs.pop("type", None) BaseNode.__init__( @@ -216,17 +217,20 @@ def search_space( return self._search_space @search_space.setter - def search_space(self, values: Dict[str, Dict[str, Union[str, int, float, dict]]]) -> None: + def search_space(self, values: Optional[Dict[str, Dict[str, Union[str, int, float, dict]]]]) -> None: """Sets the search space for the sweep job. :param values: The search space to set. - :type values: Dict[str, Dict[str, Union[str, int, float, dict]]] + :type values: Optional[Dict[str, Dict[str, Union[str, int, float, dict]]]] """ - search_space: Dict = {} - for name, value in values.items(): - # If value is a SearchSpace object, directly pass it to job.search_space[name] - search_space[name] = self._value_type_to_class(value) if isinstance(value, dict) else value - self._search_space = search_space + if values is None: + self._search_space = None + else: + search_space: Dict = {} + for name, value in values.items(): + # If value is a SearchSpace object, directly pass it to job.search_space[name] + search_space[name] = self._value_type_to_class(value) if isinstance(value, dict) else value + self._search_space = search_space @classmethod def _value_type_to_class(cls, value: Any) -> Dict: diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_component/command_component.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_component/command_component.py index 9bdcd3d1e656..72c8c506e5e9 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_component/command_component.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_component/command_component.py @@ -188,11 +188,11 @@ def instance_count(self) -> Optional[int]: return self.resources.instance_count if self.resources and not isinstance(self.resources, dict) else None @instance_count.setter - def instance_count(self, value: int) -> None: + def instance_count(self, value: Optional[int]) -> None: """Sets the number of instances or nodes to be used by the compute target. :param value: The number of instances of nodes to be used by the compute target. Defaults to 1. - :type value: int + :type value: Optional[int] """ if not value: return diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_component/component.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_component/component.py index c02a3a33794a..1cfdb45cf2ac 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_component/component.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_component/component.py @@ -230,11 +230,11 @@ def version(self) -> Optional[str]: return self._version @version.setter - def version(self, value: str) -> None: + def version(self, value: Optional[str]) -> None: """Set the version of the component. :param value: The version of the component. - :type value: str + :type value: Optional[str] """ if value: if not isinstance(value, str): diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_component/parallel_component.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_component/parallel_component.py index 3f29b1e163ed..45db26fac4b2 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_component/parallel_component.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_component/parallel_component.py @@ -179,11 +179,11 @@ def instance_count(self) -> Optional[int]: return self.resources.instance_count if self.resources and not isinstance(self.resources, dict) else None @instance_count.setter - def instance_count(self, value: int) -> None: + def instance_count(self, value: Optional[int]) -> None: """Set the value of the promoted property resources.instance_count. :param value: The value to set for resources.instance_count. - :type value: int + :type value: Optional[int] """ if not value: return @@ -204,11 +204,11 @@ def code(self) -> Optional[str]: return self.task.code if self.task else None @code.setter - def code(self, value: str) -> None: + def code(self, value: Optional[str]) -> None: """Set the value of the promoted property task.code. :param value: The value to set for task.code. - :type value: str + :type value: Optional[str] """ if not value: return diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_indexes/data_index_func.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_indexes/data_index_func.py index 884faf823f69..01b89f9a2252 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_indexes/data_index_func.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_indexes/data_index_func.py @@ -9,16 +9,15 @@ from typing import Any, Dict, Optional, Union from azure.ai.ml._utils._experimental import experimental -from azure.ai.ml.constants._common import AssetTypes, LegacyAssetTypes +from azure.ai.ml.constants._common import AssetTypes, DataIndexTypes, LegacyAssetTypes +from azure.ai.ml.constants._component import LLMRAGComponentUri from azure.ai.ml.entities import PipelineJob from azure.ai.ml.entities._builders.base_node import pipeline_node_decorator from azure.ai.ml.entities._credentials import ManagedIdentityConfiguration, UserIdentityConfiguration +from azure.ai.ml.entities._indexes.entities.data_index import DataIndex from azure.ai.ml.entities._inputs_outputs import Input, Output from azure.ai.ml.entities._job.pipeline._io import PipelineInput from azure.ai.ml.exceptions import ErrorCategory, ErrorTarget, ValidationErrorType, ValidationException -from azure.ai.ml.constants._common import DataIndexTypes -from azure.ai.ml.constants._component import LLMRAGComponentUri -from azure.ai.ml.entities._indexes.entities.data_index import DataIndex SUPPORTED_INPUTS = [ LegacyAssetTypes.PATH, @@ -745,4 +744,4 @@ def _resolve_connection_id(ml_client, connection: Optional[str] = None) -> Optio connection = ml_client.connections.get(connection_name) if connection is None: return None - return connection.id # type: ignore [attr-defined] + return connection.id # type: ignore[attr-defined, union-attr] diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_inputs_outputs/external_data.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_inputs_outputs/external_data.py index 8a4fe21f3faf..9ba8934403fc 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_inputs_outputs/external_data.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_inputs_outputs/external_data.py @@ -72,7 +72,7 @@ def __init__( ) -> None: # As an annotation, it is not allowed to initialize the name. # The name will be updated by the annotated variable name. - self.name = None + self.name: Optional[str] = None self.type = ExternalDataType.DATABASE self.connection = connection self.query = query diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/job_limits.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/job_limits.py index 7aae92635736..0797eed0b826 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/job_limits.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/job_limits.py @@ -138,11 +138,11 @@ def trial_timeout(self) -> Optional[Union[int, str]]: return self._trial_timeout @trial_timeout.setter - def trial_timeout(self, value: int) -> None: + def trial_timeout(self, value: Optional[Union[int, str]]) -> None: """Sets the timeout value for each Sweep Job trial. :param value: The timeout value, in seconds, for each Sweep Job trial. - :type value: int + :type value: Optional[Union[int, str]] """ self._trial_timeout = _get_floored_timeout(value) diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/job_resource_configuration.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/job_resource_configuration.py index fc016d59f9c2..ca6d875f8c1d 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/job_resource_configuration.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/job_resource_configuration.py @@ -137,7 +137,7 @@ def __init__( self.shm_size = shm_size self.max_instance_count = max_instance_count self.docker_args = docker_args - self._properties = None + self._properties: Optional[Properties] = None self.properties = properties @property @@ -149,11 +149,11 @@ def properties(self) -> Optional[Union[Properties, Dict]]: return self._properties @properties.setter - def properties(self, properties: Dict[str, Any]) -> None: + def properties(self, properties: Optional[Union[Properties, Dict[str, Any]]]) -> None: """Sets the properties of the job. :param properties: A dictionary of properties for the job. - :type properties: Dict[str, Any] + :type properties: Optional[Union[Properties, Dict[str, Any]]] :raises TypeError: Raised if properties is not a dictionary type. """ if properties is None: diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/parameterized_command.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/parameterized_command.py index 122d99bd09b1..15484fdf5d6f 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/parameterized_command.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/parameterized_command.py @@ -115,7 +115,19 @@ def distribution( return self._distribution @distribution.setter - def distribution(self, value: Union[dict, PyTorchDistribution, MpiDistribution]) -> None: + def distribution( + self, + value: Optional[ + Union[ + dict, + MpiDistribution, + TensorFlowDistribution, + PyTorchDistribution, + RayDistribution, + DistributionConfiguration, + ] + ], + ) -> None: """Sets the configuration for the distributed command component or job. :param value: The distribution configuration for distributed jobs. diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/pipeline/_io/base.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/pipeline/_io/base.py index b17972aeefaf..467dd1517189 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/pipeline/_io/base.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/pipeline/_io/base.py @@ -250,7 +250,7 @@ def __hash__(self) -> int: def _get_mode( cls, original_data: Optional[Union[int, bool, float, str, Input, Output, "PipelineInput"]], - data: Optional[Union[int, bool, float, str, Input, Output]], + data: Optional[Union[int, bool, float, str, Input, Output, "InputOutputBase"]], kwargs: dict, ) -> Optional[str]: """Get mode of this input/output builder. @@ -263,7 +263,8 @@ def _get_mode( :param data: Built input/output data. :type data: Union[None, int, bool, float, str azure.ai.ml.Input, - azure.ai.ml.Output] + azure.ai.ml.Output, + azure.ai.ml.entities._job.pipeline._io.InputOutputBase] :param kwargs: The kwargs :type kwargs: Dict :return: The mode diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/pipeline/pipeline_job.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/pipeline/pipeline_job.py index 7ddbbc465a22..8f4e3caf6a22 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/pipeline/pipeline_job.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/pipeline/pipeline_job.py @@ -229,11 +229,11 @@ def settings(self) -> Optional[PipelineJobSettings]: return res @settings.setter - def settings(self, value: Union[Dict, PipelineJobSettings]) -> None: + def settings(self, value: Optional[Union[Dict, PipelineJobSettings]]) -> None: """Set the pipeline job settings. :param value: The pipeline job settings. - :type value: Union[dict, ~azure.ai.ml.entities.PipelineJobSettings] + :type value: Optional[Union[dict, ~azure.ai.ml.entities.PipelineJobSettings]] """ if value is not None: if isinstance(value, PipelineJobSettings): diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/spark_job_entry_mixin.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/spark_job_entry_mixin.py index 2a1ff54978cb..3f0498bfed33 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/spark_job_entry_mixin.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/spark_job_entry_mixin.py @@ -20,7 +20,7 @@ class SparkJobEntryMixin: ) def __init__(self, **kwargs: Any): - self._entry = None + self._entry: Optional[Union[Dict[str, str], SparkJobEntry]] = None self.entry = kwargs.get("entry", None) @property diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_util.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_util.py index a3cf11eda34a..d404bc039dd9 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_util.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_util.py @@ -258,11 +258,11 @@ def validate_attribute_type(attrs_to_check: Dict[str, Any], attr_type_map: Dict[ ) -def is_empty_target(obj: Optional[Dict]) -> bool: +def is_empty_target(obj: Union[Dict, List, None]) -> bool: """Determines if it's empty target :param obj: The object to check - :type obj: Optional[Dict] + :type obj: Union[Dict, List, None] :return: True if obj is None or an empty Dict :rtype: bool """ @@ -440,11 +440,11 @@ def __init__(self, a, b): raise ValueError("Unexpected type {}".format(type(rest_dict))) -def extract_label(input_str: str) -> Union[Tuple, List]: +def extract_label(input_str: Optional[str]) -> Union[Tuple, List]: """Extract label from input string. :param input_str: The input string - :type input_str: str + :type input_str: Optional[str] :return: The rest of the string and the label :rtype: Tuple[str, Optional[str]] """ diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_data_operations.py b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_data_operations.py index d0c4507ebaab..b4820a3e1790 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_data_operations.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_data_operations.py @@ -863,7 +863,7 @@ def _set_registry_client(self, registry_name: str) -> Generator: self._operation = _client.data_versions yield finally: - self._operation_scope.registry_name = registry_ + self._operation_scope.registry_name = registry_ # type: ignore[assignment] self._operation_scope._resource_group_name = rg_ self._operation_scope._subscription_id = sub_ self._service_client = client_ diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_environment_operations.py b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_environment_operations.py index 048163d1b240..eddebef77e78 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_environment_operations.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_environment_operations.py @@ -558,7 +558,7 @@ def _set_registry_client(self, registry_name: str) -> Generator: self._version_operations = _client.environment_versions yield finally: - self._operation_scope.registry_name = registry_ + self._operation_scope.registry_name = registry_ # type: ignore[assignment] self._operation_scope._resource_group_name = rg_ self._operation_scope._subscription_id = sub_ self._service_client = client_ diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_feature_store_operations.py b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_feature_store_operations.py index a7f8e93daecc..f6152bdf8266 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_feature_store_operations.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_feature_store_operations.py @@ -491,9 +491,12 @@ def deserialize_callback(rest_obj: Any) -> FeatureStore: update_online_store_role_assignment=update_online_store_role_assignment, materialization_identity_id=( materialization_identity.resource_id - if update_workspace_role_assignment - or update_offline_store_role_assignment - or update_online_store_role_assignment + if materialization_identity + and ( + update_workspace_role_assignment + or update_offline_store_role_assignment + or update_online_store_role_assignment + ) else None ), offline_store_target=offline_store_target_to_update if update_offline_store_role_assignment else None, @@ -558,7 +561,9 @@ def begin_provision_network( module_logger.info("Provision network request initiated for feature store: %s\n", workspace_name) return poller - def _validate_offline_store(self, offline_store: MaterializationStore) -> None: + def _validate_offline_store(self, offline_store: Optional[MaterializationStore]) -> None: + if offline_store is None: + return store_regex = re.compile(STORE_REGEX_PATTERN) if offline_store and store_regex.match(offline_store.target) is None: raise ValidationError(f"Invalid AzureML offlinestore target ARM Id {offline_store.target}") diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_job_operations.py b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_job_operations.py index d5971f31d5d0..6020bd5857aa 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_job_operations.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_job_operations.py @@ -1648,17 +1648,19 @@ def _append_tid_to_studio_url(self, job: Job) -> None: try: if job.services is not None: studio_endpoint = job.services.get("Studio", None) - studio_url = studio_endpoint.endpoint - default_scopes = _resource_to_scopes(_get_base_url_from_metadata()) - module_logger.debug("default_scopes used: `%s`\n", default_scopes) - # Extract the tenant id from the credential using PyJWT - decode = jwt.decode( - self._credential.get_token(*default_scopes).token, - options={"verify_signature": False, "verify_aud": False}, - ) - tid = decode["tid"] - formatted_tid = TID_FMT.format(tid) - studio_endpoint.endpoint = studio_url + formatted_tid + if studio_endpoint is not None: + studio_url = studio_endpoint.endpoint + if studio_url is not None: + default_scopes = _resource_to_scopes(_get_base_url_from_metadata()) + module_logger.debug("default_scopes used: `%s`\n", default_scopes) + # Extract the tenant id from the credential using PyJWT + decode = jwt.decode( + self._credential.get_token(*default_scopes).token, + options={"verify_signature": False, "verify_aud": False}, + ) + tid = decode["tid"] + formatted_tid = TID_FMT.format(tid) + studio_endpoint.endpoint = studio_url + formatted_tid except Exception: # pylint: disable=W0718 module_logger.info("Proceeding with no tenant id appended to studio URL\n") diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_local_job_invoker.py b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_local_job_invoker.py index 90914e4cebcf..49ebcf27684e 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_local_job_invoker.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_local_job_invoker.py @@ -122,7 +122,7 @@ def invoke_command(project_temp_dir: Path) -> None: def get_execution_service_response( job_definition: JobBaseData, token: str, requests_pipeline: HttpPipeline -) -> Tuple[Dict[str, str], str]: +) -> Tuple[Any, Optional[str]]: """Get zip file containing local run information from Execution Service. MFE will send down a mock job contract, with service 'local'. @@ -137,7 +137,7 @@ def get_execution_service_response( :param requests_pipeline: The HttpPipeline to use when sending network requests :type requests_pipeline: HttpPipeline :return: Execution service response and snapshot ID - :rtype: Tuple[Dict[str, str], str] + :rtype: Tuple[Any, Optional[str]] """ try: local = job_definition.properties.services.get("Local", None) @@ -392,7 +392,7 @@ def start_run_if_local( credential: TokenCredential, ws_base_url: str, requests_pipeline: HttpPipeline, -) -> str: +) -> Optional[str]: """Request execution bundle from ES and run job. If Linux or WSL environment, unzip and invoke job using job spec and bootstrapper. Otherwise, invoke command locally. @@ -405,7 +405,7 @@ def start_run_if_local( :param requests_pipeline: The HttpPipeline to use when sending network requests :type requests_pipeline: HttpPipeline :return: snapshot ID - :rtype: str + :rtype: Optional[str] """ token = credential.get_token(ws_base_url + "/.default").token (zip_content, snapshot_id) = get_execution_service_response(job_definition, token, requests_pipeline) diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_model_operations.py b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_model_operations.py index 22c01bd4a3bf..3d493ab32a03 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_model_operations.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_model_operations.py @@ -22,7 +22,6 @@ CHANGED_ASSET_PATH_MSG_NO_PERSONAL_DATA, ) from azure.ai.ml._exception_helper import log_and_raise_error - from azure.ai.ml._restclient.model_dataplane import AzureMachineLearningWorkspaces as ServiceClientModelDataPlane from azure.ai.ml._restclient.v2021_10_01_dataplanepreview import ( AzureMachineLearningWorkspaces as ServiceClient102021Dataplane, @@ -696,7 +695,7 @@ def _set_registry_client(self, registry_name: str) -> Generator: self._model_versions_operation = _client.model_versions yield finally: - self._operation_scope.registry_name = registry_ + self._operation_scope.registry_name = registry_ # type: ignore[assignment] self._operation_scope._resource_group_name = rg_ self._operation_scope._subscription_id = sub_ self._service_client = client_