Add projects and hubs to AML SDK. (#35142)

* initial dump

* run black

* black, mypy, lint

* test corrections 1

* rerun black

* remove outdated test

* refactor polymorphism to not need more load functions

* syntax checks

* refactor constant to not be enum to avoid type problems

* update assets

* more test corrections

* more test corrections, feature store kind refactor

* rename type to kind

* show associated workspaces in yaml view

* assign project values from hub during creation

* formatting

* record pna in projects

* rename remove project from default_project_resource_group

* Update sdk/ml/azure-ai-ml/azure/ai/ml/constants/_common.py

Co-authored-by: kdestin <[email protected]>

* Update sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_operations.py

Co-authored-by: kdestin <[email protected]>

* Update sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_operations.py

Co-authored-by: kdestin <[email protected]>

* Update sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_operations.py

Co-authored-by: kdestin <[email protected]>

* name consistency

* undo enumification of workspace kind

---------

Co-authored-by: kdestin <[email protected]>

Author: MilesHolland

sdk/ml/azure-ai-ml/CHANGELOG.md

@@ -4,10 +4,19 @@
 
 ### Features Added
 
+- Renamed WorkspaceHub class as Hub.
+- Added Project entity class and YAML support.
+- Project and Hub operations supported by workspace operations.
+- workspace list operation supports type filtering.
+
 ### Bugs Fixed
 
 ### Breaking Changes
 
+- Removed WorkspaceHubConfig entity, and renamed WorkspaceHub to Hub.
+- workspace_hub input of Workspace class hidden, renamed to hub_id, and re-surfaced in child class Project.
+- Removed Workspace Hub Operations from ML Client.
+
 ### Other Changes
 
 ## 1.15.0 (2024-03-26)

sdk/ml/azure-ai-ml/assets.json

@@ -2,5 +2,5 @@
   "AssetsRepo": "Azure/azure-sdk-assets",
   "AssetsRepoPrefixPath": "python",
   "TagPrefix": "python/ml/azure-ai-ml",
-  "Tag": "python/ml/azure-ai-ml_baea630318"
+  "Tag": "python/ml/azure-ai-ml_cfba1d7771"
 }

sdk/ml/azure-ai-ml/azure/ai/ml/__init__.py

@@ -35,7 +35,6 @@
     load_registry,
     load_workspace,
     load_workspace_connection,
-    load_workspace_hub,
 )
 
 module_logger = logging.getLogger(__name__)
@@ -69,7 +68,6 @@
     "load_workspace",
     "load_registry",
     "load_workspace_connection",
-    "load_workspace_hub",
     "load_model_package",
 ]
 

sdk/ml/azure-ai-ml/azure/ai/ml/_ml_client.py

@@ -43,7 +43,6 @@
 from azure.ai.ml._scope_dependent_operations import OperationConfig, OperationsContainer, OperationScope
 from azure.ai.ml._telemetry.logging_handler import get_appinsights_log_handler
 from azure.ai.ml._user_agent import USER_AGENT
-from azure.ai.ml._utils._experimental import experimental
 from azure.ai.ml._utils._http_utils import HttpPipeline
 from azure.ai.ml._utils._preflight_utils import get_deployments_operation
 from azure.ai.ml._utils._registry_utils import get_registry_client
@@ -82,7 +81,6 @@
     OnlineEndpointOperations,
     RegistryOperations,
     WorkspaceConnectionsOperations,
-    WorkspaceHubOperations,
     WorkspaceOperations,
 )
 from azure.ai.ml.operations._code_operations import CodeOperations
@@ -658,15 +656,6 @@ def __init__(
             **ops_kwargs,  # type: ignore[arg-type]
         )
 
-        self._workspace_hubs = WorkspaceHubOperations(
-            self._operation_scope,
-            self._service_client_08_2023_preview,
-            self._operation_container,
-            self._credential,
-            **app_insights_handler_kwargs,  # type: ignore[arg-type]
-        )
-        self._operation_container.add(AzureMLResourceType.WORKSPACE_HUB, self._workspace_hubs)  # type: ignore[arg-type]
-
         self._operation_container.add(AzureMLResourceType.FEATURE_STORE, self._featurestores)  # type: ignore[arg-type]
         self._operation_container.add(AzureMLResourceType.FEATURE_SET, self._featuresets)
         self._operation_container.add(AzureMLResourceType.FEATURE_STORE_ENTITY, self._featurestoreentities)
@@ -811,7 +800,8 @@ def _ml_client_cli(cls, credentials: TokenCredential, subscription_id: Optional[
 
     @property
     def workspaces(self) -> WorkspaceOperations:
-        """A collection of workspace-related operations.
+        """A collection of workspace-related operations. Also manages workspace
+        sub-classes like projects and hubs.
 
         :return: Workspace operations
         :rtype: ~azure.ai.ml.operations.WorkspaceOperations
@@ -854,16 +844,6 @@ def feature_sets(self) -> FeatureSetOperations:
         """
         return self._featuresets
 
-    @property
-    @experimental
-    def workspace_hubs(self) -> WorkspaceHubOperations:
-        """A collection of workspace hub-related operations.
-
-        :return: Hub Operations
-        :rtype: HubOperations
-        """
-        return self._workspace_hubs
-
     @property
     def feature_store_entities(self) -> FeatureStoreEntityOperations:
         """A collection of feature store entity related operations.

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/_workspace_hub/workspace_hub.py

@@ -5,5 +5,7 @@
 __path__ = __import__("pkgutil").extend_path(__path__, __name__)  # type: ignore
 
 from .workspace import WorkspaceSchema
+from .ai_workspaces.project import ProjectSchema
+from .ai_workspaces.hub import HubSchema
 
-__all__ = ["WorkspaceSchema"]
+__all__ = ["WorkspaceSchema", "ProjectSchema", "HubSchema"]

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/_workspace_hub/workspace_hub_config.py

@@ -3,7 +3,3 @@
 # ---------------------------------------------------------
 
 __path__ = __import__("pkgutil").extend_path(__path__, __name__)  # type: ignore
-
-from .workspace_hub import WorkspaceHubSchema
-
-__all__ = ["WorkspaceHubSchema"]

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/workspace/__init__.py

@@ -0,0 +1,18 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+from marshmallow import fields
+
+from azure.ai.ml._schema import StringTransformedEnum
+from azure.ai.ml._schema.workspace import WorkspaceSchema
+from azure.ai.ml._utils._experimental import experimental
+from azure.ai.ml.constants import WorkspaceKind
+
+
+@experimental
+class HubSchema(WorkspaceSchema):
+    # additional_workspace_storage_accounts This field exists in the API, but is unused, and thus not surfaced yet.
+    kind = StringTransformedEnum(required=True, allowed_values=WorkspaceKind.HUB)
+    default_resource_group = fields.Str(required=False)
+    associated_workspaces = fields.List(fields.Str(), required=False, dump_only=True)

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/_workspace_hub/__init__.py renamed to sdk/ml/azure-ai-ml/azure/ai/ml/_schema/workspace/ai_workspaces/__init__.py

@@ -0,0 +1,16 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+from marshmallow import fields
+
+from azure.ai.ml._schema import StringTransformedEnum
+from azure.ai.ml._utils._experimental import experimental
+from azure.ai.ml._schema.workspace import WorkspaceSchema
+from azure.ai.ml.constants import WorkspaceKind
+
+
+@experimental
+class ProjectSchema(WorkspaceSchema):
+    kind = StringTransformedEnum(required=True, allowed_values=WorkspaceKind.PROJECT)
+    hub_id = fields.Str(required=True)

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/workspace/ai_workspaces/hub.py

@@ -8,7 +8,7 @@
 from azure.ai.ml._restclient.v2023_10_01.models import ListViewType
 
 from ._assets import IPProtectionLevel
-from ._common import AssetTypes, InputOutputModes, ModelType, Scope, TimeZone
+from ._common import AssetTypes, InputOutputModes, ModelType, Scope, TimeZone, WorkspaceKind
 from ._component import ParallelTaskType
 from ._deployment import BatchDeploymentOutputAction
 from ._job import (
@@ -63,4 +63,5 @@
     "MonitorTargetTasks",
     "IPProtectionLevel",
     "ListViewType",
+    "WorkspaceKind",
 ]

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/workspace/ai_workspaces/project.py

@@ -246,8 +246,10 @@ class AzureMLResourceType:
     """Feature store entity resource type."""
     FEATURE_STORE = "feature_store"
     """Feature store resource type."""
-    WORKSPACE_HUB = "workspace_hub"
-    """WorkspaceHub resource type."""
+    HUB = "hub"
+    """Hub resource type."""
+    PROJECT = "project"
+    """Project resource type."""
 
     NAMED_TYPES = {
         JOB,
@@ -390,6 +392,7 @@ class CommonYamlFields:
     """Name."""
     SCHEMA = "$schema"
     """Schema."""
+    KIND = "kind"
 
 
 class SchemaUrl:
@@ -490,7 +493,7 @@ class YAMLRefDocLinks:
     FEATURE_STORE = "https://aka.ms/ml-cli-v2-featurestore-yaml-reference"
     FEATURE_SET = "https://aka.ms/ml-cli-v2-featureset-yaml-reference"
     FEATURE_STORE_ENTITY = "https://aka.ms/ml-cli-v2-featurestore-entity-yaml-reference"
-    WORKSPACEHUB = "https://aka.ms/ml-cli-v2-workspace-hub-entity-yaml-reference"
+    HUB = "https://aka.ms/ml-cli-v2-workspace-hub-entity-yaml-reference"
 
 
 class YAMLRefDocSchemaNames:
@@ -910,3 +913,12 @@ class ScheduleType(str, Enum, metaclass=CaseInsensitiveEnumMeta):
 class AutoDeleteCondition(str, Enum, metaclass=CaseInsensitiveEnumMeta):
     CREATED_GREATER_THAN = "created_greater_than"
     LAST_ACCESSED_GREATER_THAN = "last_accessed_greater_than"
+
+
+class WorkspaceKind():
+    """Enum of workspace categories."""
+
+    DEFAULT = "default"
+    HUB = "hub"
+    PROJECT = "project"
+    FEATURE_STORE = "featurestore"

sdk/ml/azure-ai-ml/azure/ai/ml/constants/__init__.py

@@ -220,8 +220,9 @@
 from ._workspace.private_endpoint import EndpointConnection, PrivateEndpoint
 from ._workspace.serverless_compute import ServerlessComputeSettings
 from ._workspace.workspace import Workspace
+from ._workspace._ai_workspaces.hub import Hub
+from ._workspace._ai_workspaces.project import Project
 from ._workspace.workspace_keys import ContainerRegistryCredential, NotebookAccessKeys, WorkspaceKeys
-from ._workspace_hub.workspace_hub import WorkspaceHub, WorkspaceHubConfig
 
 __all__ = [
     "Resource",
@@ -356,8 +357,8 @@
     "AutoScaleSettings",
     "AutoPauseSettings",
     "WorkspaceModelReference",
-    "WorkspaceHub",
-    "WorkspaceHubConfig",
+    "Hub",
+    "Project",
     "Feature",
     "FeatureSet",
     "FeatureSetBackfillRequest",

sdk/ml/azure-ai-ml/azure/ai/ml/constants/_common.py

@@ -9,7 +9,6 @@
 ONLINE_MATERIALIZATION_STORE_TYPE = "redis"
 ONLINE_STORE_CONNECTION_CATEGORY = "Redis"
 DEFAULT_SPARK_RUNTIME_VERSION = "3.2.0"
-FEATURE_STORE_KIND = "featurestore"
 STORE_REGEX_PATTERN = (
     "^/?subscriptions/([^/]+)/resourceGroups/([^/]+)/providers/Microsoft.Storage"
     "/storageAccounts/([^/]+)/blobServices/default/containers/([^/]+)"

sdk/ml/azure-ai-ml/azure/ai/ml/entities/__init__.py

@@ -19,8 +19,8 @@
 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 ._constants import DEFAULT_SPARK_RUNTIME_VERSION, FEATURE_STORE_KIND
+from azure.ai.ml.constants._common import WorkspaceKind
+from ._constants import DEFAULT_SPARK_RUNTIME_VERSION
 from .materialization_store import MaterializationStore
 
 
@@ -136,7 +136,7 @@ def __init__(
             name=name,
             description=description,
             tags=tags,
-            kind=FEATURE_STORE_KIND,
+            kind=WorkspaceKind.FEATURE_STORE,
             display_name=display_name,
             location=location,
             resource_group=resource_group,

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_feature_store/_constants.py

@@ -41,7 +41,6 @@
 from azure.ai.ml.entities._validation import PathAwareSchemaValidatableMixin, ValidationResultBuilder
 from azure.ai.ml.entities._workspace.connections.workspace_connection import WorkspaceConnection
 from azure.ai.ml.entities._workspace.workspace import Workspace
-from azure.ai.ml.entities._workspace_hub.workspace_hub import WorkspaceHub
 from azure.ai.ml.exceptions import ErrorCategory, ErrorTarget, ValidationErrorType, ValidationException
 
 module_logger = logging.getLogger(__name__)
@@ -202,7 +201,8 @@ def load_workspace(
     relative_origin: Optional[str] = None,
     **kwargs: Any,
 ) -> Workspace:
-    """Load a workspace object from a yaml file.
+    """Load a workspace object from a yaml file. This includes workspace sub-classes
+    like hubs and projects.
 
     :param source: The local yaml source of a workspace. Must be either a
         path to a local file, or an already-open file.
@@ -862,45 +862,6 @@ def load_feature_store_entity(
     return cast(FeatureStoreEntity, load_common(FeatureStoreEntity, source, relative_origin, **kwargs))
 
 
-@experimental
-def load_workspace_hub(
-    source: Union[str, PathLike, IO[AnyStr]],
-    *,
-    relative_origin: Optional[str] = None,
-    **kwargs: Any,
-) -> WorkspaceHub:
-    """Load a WorkspaceHub object from a yaml file.
-
-    :param source: The local yaml source of a WorkspaceHub. Must be either a
-        path to a local file, or an already-open file.
-        If the source is a path, it will be open and read.
-        An exception is raised if the file does not exist.
-        If the source is an open file, the file will be read directly,
-        and an exception is raised if the file is not readable.
-    :type source: Union[PathLike, str, io.TextIOWrapper]
-    :keyword relative_origin: The origin to be used when deducing
-        the relative locations of files referenced in the parsed yaml.
-        Defaults to the inputted source's directory if it is a file or file path input.
-        Defaults to "./" if the source is a stream input with no name value.
-    :paramtype relative_origin: str
-    :keyword params_override: Fields to overwrite on top of the yaml file.
-        Format is [{"field1": "value1"}, {"field2": "value2"}]
-    :paramtype params_override: List[Dict]
-    :return: Loaded WorkspaceHub object.
-    :rtype: WorkspaceHub
-
-    .. admonition:: Example:
-
-        .. literalinclude:: ../samples/ml_samples_workspace.py
-            :start-after: [START load_workspace_hub]
-            :end-before: [END load_workspace_hub]
-            :language: python
-            :dedent: 8
-            :caption: Loading a Workspace Hub from a YAML config file.
-    """
-    return cast(WorkspaceHub, load_common(WorkspaceHub, source, relative_origin, **kwargs))
-
-
 @experimental
 def load_model_package(
     source: Union[str, PathLike, IO[AnyStr]],