diff --git a/sdk/ml/azure-ai-ml/.pre-commit-config.yaml b/sdk/ml/azure-ai-ml/.pre-commit-config.yaml index a41abc259c6e..10296515269a 100644 --- a/sdk/ml/azure-ai-ml/.pre-commit-config.yaml +++ b/sdk/ml/azure-ai-ml/.pre-commit-config.yaml @@ -41,7 +41,7 @@ repos: "https://pkgs.dev.azure.com/azure-sdk/public/_packaging/azure-sdk-for-python/pypi/simple/", ], ), - ("pylint", "3.0.3", []), + ("pylint", "3.2.7", []), ] # Make sure that correct versions are installed diff --git a/sdk/ml/azure-ai-ml/CHANGELOG.md b/sdk/ml/azure-ai-ml/CHANGELOG.md index 9dcf82dbab79..b1d0a36c44d4 100644 --- a/sdk/ml/azure-ai-ml/CHANGELOG.md +++ b/sdk/ml/azure-ai-ml/CHANGELOG.md @@ -3,6 +3,8 @@ ## 1.22.0 (unreleased) ### Features Added + - Added support to select firewall sku to used for provisioning azure firewall when FQDN rules are added in + AllowOnlyApprovedOutbound mode. FirewallSku options are `Standard` or `Basic`, defaults to `Standard` - Update TLS version from 1.0 to 1.2 - Added support for Distillation jobs. Can be created by importing `disillation` from `azure.ai.ml.distillation` ### Bugs Fixed diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/_ml_client.py b/sdk/ml/azure-ai-ml/azure/ai/ml/_ml_client.py index 2cc4c2e49093..1ff6643f3b60 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/_ml_client.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/_ml_client.py @@ -39,6 +39,7 @@ from azure.ai.ml._restclient.v2024_01_01_preview import AzureMachineLearningWorkspaces as ServiceClient012024Preview from azure.ai.ml._restclient.v2024_04_01_preview import AzureMachineLearningWorkspaces as ServiceClient042024Preview from azure.ai.ml._restclient.v2024_07_01_preview import AzureMachineLearningWorkspaces as ServiceClient072024Preview +from azure.ai.ml._restclient.v2024_10_01_preview import AzureMachineLearningWorkspaces as ServiceClient102024Preview from azure.ai.ml._restclient.workspace_dataplane import ( AzureMachineLearningWorkspaces as ServiceClientWorkspaceDataplane, ) @@ -381,6 +382,17 @@ def __init__( **kwargs, ) + self._service_client_10_2024_preview = ServiceClient102024Preview( + credential=self._credential, + subscription_id=( + self._ws_operation_scope._subscription_id + if registry_reference + else self._operation_scope._subscription_id + ), + base_url=base_url, + **kwargs, + ) + # A general purpose, user-configurable pipeline for making # http requests self._requests_pipeline = HttpPipeline(**kwargs) @@ -478,7 +490,7 @@ def __init__( self._workspaces = WorkspaceOperations( self._ws_operation_scope if registry_reference else self._operation_scope, - self._service_client_07_2024_preview, + self._service_client_10_2024_preview, self._operation_container, self._credential, requests_pipeline=self._requests_pipeline, @@ -489,7 +501,7 @@ def __init__( self._workspace_outbound_rules = WorkspaceOutboundRuleOperations( self._operation_scope, - self._service_client_07_2024_preview, + self._service_client_10_2024_preview, self._operation_container, self._credential, **kwargs, @@ -706,7 +718,7 @@ def __init__( self._featurestores = FeatureStoreOperations( self._operation_scope, - self._service_client_07_2024_preview, + self._service_client_10_2024_preview, self._operation_container, self._credential, **app_insights_handler_kwargs, # type: ignore[arg-type] diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/_schema/workspace/networking.py b/sdk/ml/azure-ai-ml/azure/ai/ml/_schema/workspace/networking.py index feee400dd8f5..9c805b783e80 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/_schema/workspace/networking.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/_schema/workspace/networking.py @@ -4,18 +4,19 @@ # pylint: disable=unused-argument,no-else-return -from marshmallow import fields, EXCLUDE +from marshmallow import EXCLUDE, fields from marshmallow.decorators import post_load, pre_dump + +from azure.ai.ml._schema.core.fields import NestedField, StringTransformedEnum, UnionField from azure.ai.ml._schema.core.schema_meta import PatchedSchemaMeta -from azure.ai.ml._schema.core.fields import StringTransformedEnum, NestedField, UnionField +from azure.ai.ml._utils.utils import _snake_to_camel, camel_to_snake +from azure.ai.ml.constants._workspace import FirewallSku, IsolationMode, OutboundRuleCategory from azure.ai.ml.entities._workspace.networking import ( - ManagedNetwork, FqdnDestination, - ServiceTagDestination, + ManagedNetwork, PrivateEndpointDestination, + ServiceTagDestination, ) -from azure.ai.ml.constants._workspace import IsolationMode, OutboundRuleCategory -from azure.ai.ml._utils.utils import camel_to_snake, _snake_to_camel class ManagedNetworkStatusSchema(metaclass=PatchedSchemaMeta): @@ -184,13 +185,31 @@ class ManagedNetworkSchema(metaclass=PatchedSchemaMeta): ), allow_none=True, ) + firewall_sku = StringTransformedEnum( + allowed_values=[ + FirewallSku.STANDARD, + FirewallSku.BASIC, + ], + casing_transform=camel_to_snake, + metadata={"description": "Firewall sku for FQDN rules in AllowOnlyApprovedOutbound mode"}, + ) network_id = fields.Str(required=False, dump_only=True) status = NestedField(ManagedNetworkStatusSchema, allow_none=False, unknown=EXCLUDE) @post_load def make(self, data, **kwargs): outbound_rules = data.get("outbound_rules", False) + + firewall_sku = data.get("firewall_sku", False) + firewall_sku_value = _snake_to_camel(data["firewall_sku"]) if firewall_sku else FirewallSku.STANDARD + if outbound_rules: - return ManagedNetwork(isolation_mode=_snake_to_camel(data["isolation_mode"]), outbound_rules=outbound_rules) + return ManagedNetwork( + isolation_mode=_snake_to_camel(data["isolation_mode"]), + outbound_rules=outbound_rules, + firewall_sku=firewall_sku_value, + ) else: - return ManagedNetwork(isolation_mode=_snake_to_camel(data["isolation_mode"])) + return ManagedNetwork( + isolation_mode=_snake_to_camel(data["isolation_mode"]), firewall_sku=firewall_sku_value + ) diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/constants/_workspace.py b/sdk/ml/azure-ai-ml/azure/ai/ml/constants/_workspace.py index caeed99bbe24..a47c4e5b0b4a 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/constants/_workspace.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/constants/_workspace.py @@ -23,6 +23,13 @@ class IsolationMode: ALLOW_ONLY_APPROVED_OUTBOUND = "AllowOnlyApprovedOutbound" +class FirewallSku: + """Firewall Sku for FQDN rules in AllowOnlyApprovedOutbound.""" + + STANDARD = "Standard" + BASIC = "Basic" + + class OutboundRuleCategory: """Category for a managed network outbound rule.""" diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_feature_store/feature_store.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_feature_store/feature_store.py index dde51afd2d7e..87ea9e70bbce 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_feature_store/feature_store.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_feature_store/feature_store.py @@ -9,9 +9,9 @@ from pathlib import Path from typing import Any, Dict, Optional, Union -from azure.ai.ml._restclient.v2024_07_01_preview.models import Workspace as RestWorkspace +from azure.ai.ml._restclient.v2024_10_01_preview.models import Workspace as RestWorkspace from azure.ai.ml._schema._feature_store.feature_store_schema import FeatureStoreSchema -from azure.ai.ml.constants._common import BASE_PATH_CONTEXT_KEY, PARAMS_OVERRIDE_KEY +from azure.ai.ml.constants._common import BASE_PATH_CONTEXT_KEY, PARAMS_OVERRIDE_KEY, WorkspaceKind from azure.ai.ml.entities._credentials import IdentityConfiguration, ManagedIdentityConfiguration from azure.ai.ml.entities._util import load_from_dict from azure.ai.ml.entities._workspace.compute_runtime import ComputeRuntime @@ -19,7 +19,7 @@ from azure.ai.ml.entities._workspace.feature_store_settings import FeatureStoreSettings from azure.ai.ml.entities._workspace.networking import ManagedNetwork from azure.ai.ml.entities._workspace.workspace import Workspace -from azure.ai.ml.constants._common import WorkspaceKind + from ._constants import DEFAULT_SPARK_RUNTIME_VERSION from .materialization_store import MaterializationStore diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/_ai_workspaces/hub.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/_ai_workspaces/hub.py index 3bae3f8d58b5..8596f6d9e60f 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/_ai_workspaces/hub.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/_ai_workspaces/hub.py @@ -5,10 +5,8 @@ # pylint: disable=too-many-instance-attributes,protected-access from typing import Any, Dict, List, Optional -from azure.ai.ml._restclient.v2024_07_01_preview.models import ( - Workspace as RestWorkspace, - WorkspaceHubConfig as RestWorkspaceHubConfig, -) +from azure.ai.ml._restclient.v2024_10_01_preview.models import Workspace as RestWorkspace +from azure.ai.ml._restclient.v2024_10_01_preview.models import WorkspaceHubConfig as RestWorkspaceHubConfig from azure.ai.ml._schema.workspace import HubSchema from azure.ai.ml._utils._experimental import experimental from azure.ai.ml.constants._common import WorkspaceKind diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/diagnose.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/diagnose.py index 29a9759a7fc3..fa923dc46391 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/diagnose.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/diagnose.py @@ -3,13 +3,17 @@ # --------------------------------------------------------- import json -from typing import Any, Dict, Optional, List +from typing import Any, Dict, List, Optional -from azure.ai.ml._restclient.v2024_07_01_preview.models import ( +from azure.ai.ml._restclient.v2024_10_01_preview.models import ( DiagnoseRequestProperties as RestDiagnoseRequestProperties, - DiagnoseResponseResult as RestDiagnoseResponseResult, +) +from azure.ai.ml._restclient.v2024_10_01_preview.models import DiagnoseResponseResult as RestDiagnoseResponseResult +from azure.ai.ml._restclient.v2024_10_01_preview.models import ( DiagnoseResponseResultValue as RestDiagnoseResponseResultValue, - DiagnoseResult as RestDiagnoseResult, +) +from azure.ai.ml._restclient.v2024_10_01_preview.models import DiagnoseResult as RestDiagnoseResult +from azure.ai.ml._restclient.v2024_10_01_preview.models import ( DiagnoseWorkspaceParameters as RestDiagnoseWorkspaceParameters, ) diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/feature_store_settings.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/feature_store_settings.py index d680f1201a62..0b5388491c9f 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/feature_store_settings.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/feature_store_settings.py @@ -6,7 +6,7 @@ from typing import Optional -from azure.ai.ml._restclient.v2024_07_01_preview.models import FeatureStoreSettings as RestFeatureStoreSettings +from azure.ai.ml._restclient.v2024_10_01_preview.models import FeatureStoreSettings as RestFeatureStoreSettings from azure.ai.ml.entities._mixins import RestTranslatableMixin from .compute_runtime import ComputeRuntime diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/networking.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/networking.py index c216913400ff..4576eac9b54c 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/networking.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/networking.py @@ -5,15 +5,21 @@ from abc import ABC from typing import Any, Dict, List, Optional -from azure.ai.ml._restclient.v2024_07_01_preview.models import ( - FqdnOutboundRule as RestFqdnOutboundRule, +from azure.ai.ml._restclient.v2024_10_01_preview.models import FqdnOutboundRule as RestFqdnOutboundRule +from azure.ai.ml._restclient.v2024_10_01_preview.models import ( ManagedNetworkProvisionStatus as RestManagedNetworkProvisionStatus, - ManagedNetworkSettings as RestManagedNetwork, +) +from azure.ai.ml._restclient.v2024_10_01_preview.models import ManagedNetworkSettings as RestManagedNetwork +from azure.ai.ml._restclient.v2024_10_01_preview.models import ( PrivateEndpointDestination as RestPrivateEndpointOutboundRuleDestination, +) +from azure.ai.ml._restclient.v2024_10_01_preview.models import ( PrivateEndpointOutboundRule as RestPrivateEndpointOutboundRule, +) +from azure.ai.ml._restclient.v2024_10_01_preview.models import ( ServiceTagDestination as RestServiceTagOutboundRuleDestination, - ServiceTagOutboundRule as RestServiceTagOutboundRule, ) +from azure.ai.ml._restclient.v2024_10_01_preview.models import ServiceTagOutboundRule as RestServiceTagOutboundRule from azure.ai.ml.constants._workspace import IsolationMode, OutboundRuleCategory, OutboundRuleType @@ -253,6 +259,8 @@ class ManagedNetwork: :param isolation_mode: Isolation of the managed network, defaults to Disabled. :type isolation_mode: str + :param firewall_sku: Firewall Sku for FQDN rules in AllowOnlyApprovedOutbound.. + :type firewall_sku: str :param outbound_rules: List of outbound rules for the managed network. :type outbound_rules: List[~azure.ai.ml.entities.OutboundRule] :param network_id: Network id for the managed network, not meant to be set by user. @@ -271,10 +279,12 @@ def __init__( *, isolation_mode: str = IsolationMode.DISABLED, outbound_rules: Optional[List[OutboundRule]] = None, + firewall_sku: Optional[str] = None, network_id: Optional[str] = None, **kwargs: Any, ) -> None: self.isolation_mode = isolation_mode + self.firewall_sku = firewall_sku self.network_id = network_id self.outbound_rules = outbound_rules self.status = kwargs.pop("status", None) @@ -289,7 +299,9 @@ def _to_rest_object(self) -> RestManagedNetwork: if self.outbound_rules else {} ) - return RestManagedNetwork(isolation_mode=self.isolation_mode, outbound_rules=rest_outbound_rules) + return RestManagedNetwork( + isolation_mode=self.isolation_mode, outbound_rules=rest_outbound_rules, firewall_sku=self.firewall_sku + ) @classmethod def _from_rest_object(cls, obj: RestManagedNetwork) -> "ManagedNetwork": @@ -306,6 +318,7 @@ def _from_rest_object(cls, obj: RestManagedNetwork) -> "ManagedNetwork": outbound_rules=from_rest_outbound_rules, # type: ignore[arg-type] network_id=obj.network_id, status=obj.status, + firewall_sku=obj.firewall_sku, ) diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/serverless_compute.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/serverless_compute.py index b0e1c45aec0a..b78ede0609e4 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/serverless_compute.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/serverless_compute.py @@ -5,7 +5,7 @@ from marshmallow.exceptions import ValidationError -from azure.ai.ml._restclient.v2024_07_01_preview.models import ( +from azure.ai.ml._restclient.v2024_10_01_preview.models import ( ServerlessComputeSettings as RestServerlessComputeSettings, ) from azure.ai.ml._schema._utils.utils import ArmId diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/workspace.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/workspace.py index 4317cb3f338f..e07ab0c3d9ca 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/workspace.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/workspace.py @@ -8,13 +8,13 @@ from pathlib import Path from typing import IO, Any, AnyStr, Dict, List, Optional, Tuple, Type, Union -from azure.ai.ml._restclient.v2024_07_01_preview.models import ( - FeatureStoreSettings as RestFeatureStoreSettings, - ManagedNetworkSettings as RestManagedNetwork, - ManagedServiceIdentity as RestManagedServiceIdentity, +from azure.ai.ml._restclient.v2024_10_01_preview.models import FeatureStoreSettings as RestFeatureStoreSettings +from azure.ai.ml._restclient.v2024_10_01_preview.models import ManagedNetworkSettings as RestManagedNetwork +from azure.ai.ml._restclient.v2024_10_01_preview.models import ManagedServiceIdentity as RestManagedServiceIdentity +from azure.ai.ml._restclient.v2024_10_01_preview.models import ( ServerlessComputeSettings as RestServerlessComputeSettings, - Workspace as RestWorkspace, ) +from azure.ai.ml._restclient.v2024_10_01_preview.models import Workspace as RestWorkspace from azure.ai.ml._schema.workspace.workspace import WorkspaceSchema from azure.ai.ml._utils.utils import dump_yaml_to_file from azure.ai.ml.constants._common import ( diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/workspace_keys.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/workspace_keys.py index 5bf669bcbe35..4213b419dbf4 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/workspace_keys.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/workspace_keys.py @@ -4,7 +4,7 @@ from typing import List, Optional -from azure.ai.ml._restclient.v2024_07_01_preview.models import ListWorkspaceKeysResult +from azure.ai.ml._restclient.v2024_10_01_preview.models import ListWorkspaceKeysResult class ContainerRegistryCredential: 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 bde02420a16f..99b4701c331e 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 @@ -10,8 +10,8 @@ from marshmallow import ValidationError -from azure.ai.ml._restclient.v2024_07_01_preview import AzureMachineLearningWorkspaces as ServiceClient072024Preview -from azure.ai.ml._restclient.v2024_07_01_preview.models import ManagedNetworkProvisionOptions +from azure.ai.ml._restclient.v2024_10_01_preview import AzureMachineLearningWorkspaces as ServiceClient102024Preview +from azure.ai.ml._restclient.v2024_10_01_preview.models import ManagedNetworkProvisionOptions from azure.ai.ml._scope_dependent_operations import OperationsContainer, OperationScope from azure.ai.ml._telemetry import ActivityType, monitor_with_activity from azure.ai.ml._utils._logger_utils import OpsLogger @@ -19,10 +19,10 @@ from azure.ai.ml.constants import ManagedServiceIdentityType from azure.ai.ml.constants._common import Scope, WorkspaceKind from azure.ai.ml.entities import ( - WorkspaceConnection, IdentityConfiguration, ManagedIdentityConfiguration, ManagedNetworkProvisionStatus, + WorkspaceConnection, ) from azure.ai.ml.entities._feature_store._constants import ( OFFLINE_MATERIALIZATION_STORE_TYPE, @@ -58,7 +58,7 @@ class FeatureStoreOperations(WorkspaceOperationsBase): def __init__( self, operation_scope: OperationScope, - service_client: ServiceClient072024Preview, + service_client: ServiceClient102024Preview, all_operations: OperationsContainer, credentials: Optional[TokenCredential] = None, **kwargs: Dict, diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_operations.py b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_operations.py index 8ee7d1fc7ad2..b04e516891d5 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_operations.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_operations.py @@ -8,8 +8,8 @@ from marshmallow import ValidationError -from azure.ai.ml._restclient.v2024_07_01_preview import AzureMachineLearningWorkspaces as ServiceClient072024Preview -from azure.ai.ml._restclient.v2024_07_01_preview.models import ManagedNetworkProvisionOptions +from azure.ai.ml._restclient.v2024_10_01_preview import AzureMachineLearningWorkspaces as ServiceClient102024Preview +from azure.ai.ml._restclient.v2024_10_01_preview.models import ManagedNetworkProvisionOptions from azure.ai.ml._scope_dependent_operations import OperationsContainer, OperationScope from azure.ai.ml._telemetry import ActivityType, monitor_with_activity from azure.ai.ml._utils._http_utils import HttpPipeline @@ -52,7 +52,7 @@ class WorkspaceOperations(WorkspaceOperationsBase): def __init__( self, operation_scope: OperationScope, - service_client: ServiceClient072024Preview, + service_client: ServiceClient102024Preview, all_operations: OperationsContainer, credentials: Optional[TokenCredential] = None, **kwargs: Any, diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_operations_base.py b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_operations_base.py index 528c35d92180..e4f8c9157005 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_operations_base.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_operations_base.py @@ -11,8 +11,8 @@ from azure.ai.ml._arm_deployments import ArmDeploymentExecutor from azure.ai.ml._arm_deployments.arm_helper import get_template -from azure.ai.ml._restclient.v2024_07_01_preview import AzureMachineLearningWorkspaces as ServiceClient072024Preview -from azure.ai.ml._restclient.v2024_07_01_preview.models import ( +from azure.ai.ml._restclient.v2024_10_01_preview import AzureMachineLearningWorkspaces as ServiceClient102024Preview +from azure.ai.ml._restclient.v2024_10_01_preview.models import ( EncryptionKeyVaultUpdateProperties, EncryptionUpdateProperties, WorkspaceUpdateParameters, @@ -60,7 +60,7 @@ class WorkspaceOperationsBase(ABC): def __init__( self, operation_scope: OperationScope, - service_client: ServiceClient072024Preview, + service_client: ServiceClient102024Preview, all_operations: OperationsContainer, credentials: Optional[TokenCredential] = None, **kwargs: Dict, diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_outbound_rule_operations.py b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_outbound_rule_operations.py index db3c1b0e31a7..31eaf0947cb9 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_outbound_rule_operations.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_outbound_rule_operations.py @@ -4,8 +4,8 @@ from typing import Any, Dict, Iterable, Optional -from azure.ai.ml._restclient.v2024_07_01_preview import AzureMachineLearningWorkspaces as ServiceClient072024Preview -from azure.ai.ml._restclient.v2024_07_01_preview.models import OutboundRuleBasicResource +from azure.ai.ml._restclient.v2024_10_01_preview import AzureMachineLearningWorkspaces as ServiceClient102024Preview +from azure.ai.ml._restclient.v2024_10_01_preview.models import OutboundRuleBasicResource from azure.ai.ml._scope_dependent_operations import OperationsContainer, OperationScope from azure.ai.ml._telemetry import ActivityType, monitor_with_activity from azure.ai.ml._utils._logger_utils import OpsLogger @@ -28,7 +28,7 @@ class WorkspaceOutboundRuleOperations: def __init__( self, operation_scope: OperationScope, - service_client: ServiceClient072024Preview, + service_client: ServiceClient102024Preview, all_operations: OperationsContainer, credentials: TokenCredential = None, **kwargs: Dict, diff --git a/sdk/ml/azure-ai-ml/samples/ml_samples_workspace.py b/sdk/ml/azure-ai-ml/samples/ml_samples_workspace.py index 18232c26d0b1..08011726a89c 100644 --- a/sdk/ml/azure-ai-ml/samples/ml_samples_workspace.py +++ b/sdk/ml/azure-ai-ml/samples/ml_samples_workspace.py @@ -16,9 +16,11 @@ """ import os + +from ml_samples_compute import handle_resource_exists_error + from azure.ai.ml import MLClient from azure.identity import DefaultAzureCredential -from ml_samples_compute import handle_resource_exists_error subscription_id = os.environ["AZURE_SUBSCRIPTION_ID"] resource_group = os.environ["RESOURCE_GROUP_NAME"] @@ -74,13 +76,14 @@ def ml_workspace_config_sample_snippets_entities(self): # [END customermanagedkey] # [START workspace_managed_network] + from azure.ai.ml.constants._workspace import FirewallSku from azure.ai.ml.entities import ( - Workspace, - ManagedNetwork, + FqdnDestination, IsolationMode, - ServiceTagDestination, + ManagedNetwork, PrivateEndpointDestination, - FqdnDestination, + ServiceTagDestination, + Workspace, ) # Example private endpoint outbound to a blob @@ -99,9 +102,14 @@ def ml_workspace_config_sample_snippets_entities(self): # Example FQDN rule pypirule = FqdnDestination(name="pypirule", destination="pypi.org") + # Example FirewallSku + # FirewallSku is an optional parameter, when unspecified this will default to FirewallSku.Standard + firewallSku = FirewallSku.BASIC + network = ManagedNetwork( isolation_mode=IsolationMode.ALLOW_ONLY_APPROVED_OUTBOUND, outbound_rules=[blobrule, datafactoryrule, pypirule], + firewall_sku=firewallSku, ) # Workspace configuration @@ -230,8 +238,7 @@ def ml_workspace_config_sample_snippets_operations(self): # [START create_or_update_connection] from azure.ai.ml import MLClient - from azure.ai.ml.entities import WorkspaceConnection - from azure.ai.ml.entities import UsernamePasswordConfiguration + from azure.ai.ml.entities import UsernamePasswordConfiguration, WorkspaceConnection ml_client_ws = MLClient(credential, subscription_id, resource_group, workspace_name="test-ws") wps_connection = WorkspaceConnection( diff --git a/sdk/ml/azure-ai-ml/tests/test_configs/workspace/workspace_mvnet_with_firewallsku.yaml b/sdk/ml/azure-ai-ml/tests/test_configs/workspace/workspace_mvnet_with_firewallsku.yaml new file mode 100644 index 000000000000..6e7942bd57e5 --- /dev/null +++ b/sdk/ml/azure-ai-ml/tests/test_configs/workspace/workspace_mvnet_with_firewallsku.yaml @@ -0,0 +1,19 @@ +name: unittest_test_mvnet_firewallsku +location: centraluseuap +managed_network: + isolation_mode: allow_only_approved_outbound + outbound_rules: + - name: microsoft + destination: 'microsoft.com' + type: fqdn + - name: appGwRule + destination: + service_resource_id: /someappgwid + spark_enabled: false + subresource_target: appGwPrivateFrontendIpIPv4 + fqdns: ["contoso.com", "contoso2.com"] + type: private_endpoint + - name: pytorch + destination: '*.pytorch.org' + type: fqdn + firewall_sku: basic diff --git a/sdk/ml/azure-ai-ml/tests/workspace/unittests/test_workspace_entity.py b/sdk/ml/azure-ai-ml/tests/workspace/unittests/test_workspace_entity.py index 934fdcf2fbd8..c61c70231896 100644 --- a/sdk/ml/azure-ai-ml/tests/workspace/unittests/test_workspace_entity.py +++ b/sdk/ml/azure-ai-ml/tests/workspace/unittests/test_workspace_entity.py @@ -4,12 +4,9 @@ from marshmallow.exceptions import ValidationError from azure.ai.ml import load_workspace -from azure.ai.ml._restclient.v2024_07_01_preview.models import Workspace -from azure.ai.ml.constants._workspace import IsolationMode -from azure.ai.ml.entities import ( - ServerlessComputeSettings, - Workspace, -) +from azure.ai.ml._restclient.v2024_10_01_preview.models import Workspace +from azure.ai.ml.constants._workspace import FirewallSku, IsolationMode +from azure.ai.ml.entities import ServerlessComputeSettings, Workspace @pytest.mark.unittest @@ -81,6 +78,7 @@ def test_workspace_load_yamls_to_test_outbound_rule_load(self): assert workspace.managed_network is not None assert workspace.managed_network.isolation_mode == IsolationMode.ALLOW_ONLY_APPROVED_OUTBOUND + rules = workspace.managed_network.outbound_rules assert rules[0].name == "microsoft" assert rules[0].destination == "microsoft.com" @@ -98,3 +96,25 @@ def test_workspace_load_yamls_to_test_outbound_rule_load(self): assert rules[2].port_ranges == "80, 8080-8089" assert "168.63.129.16" in rules[2].address_prefixes assert "10.0.0.0/24" in rules[2].address_prefixes + + def test_workspace_load_yamls_to_test_firewallsku_load(self): + workspace = load_workspace("./tests/test_configs/workspace/workspace_mvnet_with_firewallsku.yaml") + + assert workspace.managed_network is not None + assert workspace.managed_network.isolation_mode == IsolationMode.ALLOW_ONLY_APPROVED_OUTBOUND + rules = workspace.managed_network.outbound_rules + + assert workspace.managed_network.firewall_sku == "Basic" + + assert rules[0].name == "microsoft" + assert rules[0].destination == "microsoft.com" + + assert rules[1].name == "appGwRule" + assert rules[1].service_resource_id == "/someappgwid" + assert rules[1].spark_enabled == False + assert rules[1].subresource_target == "appGwPrivateFrontendIpIPv4" + assert "contoso.com" in rules[1].fqdns + assert "contoso2.com" in rules[1].fqdns + + assert rules[2].name == "pytorch" + assert rules[2].destination == "*.pytorch.org" diff --git a/sdk/ml/azure-ai-ml/tests/workspace/unittests/test_workspace_operations.py b/sdk/ml/azure-ai-ml/tests/workspace/unittests/test_workspace_operations.py index 97e43a1cc1c5..9470ebf82cc9 100644 --- a/sdk/ml/azure-ai-ml/tests/workspace/unittests/test_workspace_operations.py +++ b/sdk/ml/azure-ai-ml/tests/workspace/unittests/test_workspace_operations.py @@ -1,3 +1,4 @@ +import urllib.parse from typing import Optional from unittest.mock import ANY, DEFAULT, MagicMock, Mock, patch from uuid import UUID, uuid4 @@ -6,10 +7,10 @@ from pytest_mock import MockFixture from azure.ai.ml import load_workspace -from azure.ai.ml._restclient.v2024_07_01_preview.models import ( +from azure.ai.ml._restclient.v2024_10_01_preview.models import ( ServerlessComputeSettings as RestServerlessComputeSettings, - Workspace as RestWorkspace, ) +from azure.ai.ml._restclient.v2024_10_01_preview.models import Workspace as RestWorkspace from azure.ai.ml._scope_dependent_operations import OperationScope from azure.ai.ml.entities import ( CustomerManagedKey, @@ -20,7 +21,6 @@ ) from azure.ai.ml.operations import WorkspaceOperations from azure.core.polling import LROPoller -import urllib.parse @pytest.fixture diff --git a/sdk/ml/azure-ai-ml/tests/workspace/unittests/test_workspace_operations_base.py b/sdk/ml/azure-ai-ml/tests/workspace/unittests/test_workspace_operations_base.py index c3dcc424e34c..3b98ff22d8c3 100644 --- a/sdk/ml/azure-ai-ml/tests/workspace/unittests/test_workspace_operations_base.py +++ b/sdk/ml/azure-ai-ml/tests/workspace/unittests/test_workspace_operations_base.py @@ -1,3 +1,4 @@ +import urllib.parse from typing import Optional from unittest.mock import DEFAULT, MagicMock, Mock, patch from uuid import UUID, uuid4 @@ -5,11 +6,11 @@ import pytest from pytest_mock import MockFixture -from azure.ai.ml._restclient.v2024_07_01_preview.models import ( +from azure.ai.ml._restclient.v2024_10_01_preview.models import ( EncryptionKeyVaultUpdateProperties, EncryptionUpdateProperties, ) -from azure.ai.ml._restclient.v2024_07_01_preview.models import ( +from azure.ai.ml._restclient.v2024_10_01_preview.models import ( ServerlessComputeSettings as RestServerlessComputeSettings, ) from azure.ai.ml._scope_dependent_operations import OperationScope @@ -31,7 +32,6 @@ ) from azure.ai.ml.operations._workspace_operations_base import WorkspaceOperationsBase from azure.core.polling import LROPoller -import urllib.parse @pytest.fixture diff --git a/sdk/ml/azure-ai-ml/tests/workspace/unittests/test_workspace_related_restclients.py b/sdk/ml/azure-ai-ml/tests/workspace/unittests/test_workspace_related_restclients.py index f90ae3fc9123..68881e13717f 100644 --- a/sdk/ml/azure-ai-ml/tests/workspace/unittests/test_workspace_related_restclients.py +++ b/sdk/ml/azure-ai-ml/tests/workspace/unittests/test_workspace_related_restclients.py @@ -1,39 +1,47 @@ from typing import Optional import pytest -from azure.ai.ml.entities import ( - Workspace, - FqdnDestination, - ServiceTagDestination, - PrivateEndpointDestination, - FeatureStore, - Hub, -) -from azure.ai.ml._restclient.v2024_07_01_preview.models import ( - Workspace as RestWorkspace, - ManagedNetworkSettings as RestManagedNetwork, - FqdnOutboundRule as RestFqdnOutboundRule, - PrivateEndpointOutboundRule as RestPrivateEndpointOutboundRule, - PrivateEndpointDestination as RestPrivateEndpointOutboundRuleDestination, - ServiceTagOutboundRule as RestServiceTagOutboundRule, - ServiceTagDestination as RestServiceTagOutboundRuleDestination, +from azure.ai.ml import MLClient +from azure.ai.ml._restclient.v2024_10_01_preview.models import FeatureStoreSettings as RestFeatureStoreSettings +from azure.ai.ml._restclient.v2024_10_01_preview.models import FqdnOutboundRule as RestFqdnOutboundRule +from azure.ai.ml._restclient.v2024_10_01_preview.models import ( ManagedNetworkProvisionStatus as RestManagedNetworkProvisionStatus, - FeatureStoreSettings as RestFeatureStoreSettings, - ManagedServiceIdentity as RestManagedServiceIdentity, +) +from azure.ai.ml._restclient.v2024_10_01_preview.models import ManagedNetworkSettings as RestManagedNetwork +from azure.ai.ml._restclient.v2024_10_01_preview.models import ManagedServiceIdentity as RestManagedServiceIdentity +from azure.ai.ml._restclient.v2024_10_01_preview.models import ( + PrivateEndpointDestination as RestPrivateEndpointOutboundRuleDestination, +) +from azure.ai.ml._restclient.v2024_10_01_preview.models import ( + PrivateEndpointOutboundRule as RestPrivateEndpointOutboundRule, +) +from azure.ai.ml._restclient.v2024_10_01_preview.models import ( ServerlessComputeSettings as RestServerlessComputeSettings, - UserAssignedIdentity, - # this one only for workspace hubs - WorkspaceHubConfig as RestWorkspaceHubConfig, ) -from azure.ai.ml._restclient.v2024_07_01_preview.operations import ( - WorkspacesOperations as RestClientWorkspacesOperations, +from azure.ai.ml._restclient.v2024_10_01_preview.models import ( + ServiceTagDestination as RestServiceTagOutboundRuleDestination, +) +from azure.ai.ml._restclient.v2024_10_01_preview.models import ServiceTagOutboundRule as RestServiceTagOutboundRule +from azure.ai.ml._restclient.v2024_10_01_preview.models import UserAssignedIdentity +from azure.ai.ml._restclient.v2024_10_01_preview.models import ( + Workspace as RestWorkspace, # this one only for workspace hubs +) +from azure.ai.ml._restclient.v2024_10_01_preview.models import WorkspaceHubConfig as RestWorkspaceHubConfig +from azure.ai.ml._restclient.v2024_10_01_preview.operations import ( ManagedNetworkSettingsRuleOperations as RestClientManagedNetworkSettingsRuleOperations, ) - +from azure.ai.ml._restclient.v2024_10_01_preview.operations import ( + WorkspacesOperations as RestClientWorkspacesOperations, +) from azure.ai.ml.constants._workspace import IsolationMode -from azure.ai.ml import ( - MLClient, +from azure.ai.ml.entities import ( + FeatureStore, + FqdnDestination, + Hub, + PrivateEndpointDestination, + ServiceTagDestination, + Workspace, ) from azure.identity import DefaultAzureCredential @@ -177,8 +185,8 @@ def test_workspace_hub_entity_from_rest_to_ensure_restclient_versions_match(self def test_feature_store_entity_from_rest_to_ensure_restclient_versions_match(self): rest_ws = get_test_rest_workspace_with_all_details() - sdk_featurestore = FeatureStore._from_rest_object(rest_ws) + assert sdk_featurestore.managed_network is not None assert sdk_featurestore.managed_network.isolation_mode == IsolationMode.ALLOW_ONLY_APPROVED_OUTBOUND rules = sdk_featurestore.managed_network.outbound_rules