sdk/python/assets/assets-in-registry/share-data-using-registry.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Share data across workspaces\n",
    "\n",
    "This is the companion notebook for the article on sharing components, environments and models across workspaces: https://learn.microsoft.com/en-us/azure/machine-learning/how-to-share-data-across-workspaces-with-registries\n",
    "\n",
    "### Prerequisites\n",
    "Review the prerequisites section in the below article to get all the details and relevant links:\n",
    "\n",
    "https://learn.microsoft.com/en-us/azure/machine-learning/how-to-share-data-across-workspaces-with-registries\n",
    "\n",
    "- Familiarity with Azure Machine Learning registries and Data concepts in Azure Machine Learning.\n",
    "- An Azure Machine Learning registry (preview) to share data.\n",
    "- An Azure Machine Learning workspace. (The Azure region (location) where you create your workspace must be in the list of supported regions for Azure Machine Learning registry.)\n",
    "- The Azure Machine Learning Python SDK v2\n",
    "\n",
    "\n",
    "### Overview and scenarios\n",
    "\n",
    "Azure Machine Learning registry (preview) enables you to collaborate across workspaces within your organization. Using registries, you can share models, components, environments and data.\n",
    "\n",
    "#### Key scenario addressed by data sharing using Azure Machine Learning registry\n",
    "\n",
    "You may want to have data shared across multiple teams, projects, or workspaces in a central location. Such data doesn't have sensitive access controls and can be broadly used in the organization.  \n",
    "\n",
    "Examples include:\n",
    "* A team wants to share a public dataset that is preprocessed and ready to use in experiments.\n",
    "* Your organization has acquired a particular dataset for a project from an external vendor and wants to make it available to all teams working on a project.\n",
    "* A team wants to share data assets across workspaces in different regions.\n",
    "\n",
    "#### Scenarios NOT addressed by data sharing using Azure Machine Learning registry\n",
    "\n",
    "* Sharing sensitive data that requires fine grained access control. You can't create a data asset in a registry to share with a small subset of users/workspaces while the registry is accessible by many other users in the org.\n",
    "* SSharing data that is available in existing storage that must not be copied or is too large or too expensive to be copied. Whenever data assets are created in a registry, a copy of data is ingested into the registry storage so that it can be replicated.\n",
    "\n",
    "\n",
    "### Goals\n",
    "\n",
    "In this article, you'll learn how to:\n",
    "* Create a data asset in the registry.\n",
    "* Use the data asset from registry as input to a model training job in a workspace.\n",
    "* Share an existing data asset from workspace to registry"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import required libraries\n",
    "from azure.identity import DefaultAzureCredential, InteractiveBrowserCredential\n",
    "\n",
    "from azure.ai.ml import MLClient, Input, Output\n",
    "from azure.ai.ml.dsl import pipeline\n",
    "from azure.ai.ml import load_component\n",
    "from azure.ai.ml.entities import (\n",
    "    Environment,\n",
    "    BuildContext,\n",
    "    Model,\n",
    "    Data,\n",
    "    CodeConfiguration,\n",
    ")\n",
    "from azure.ai.ml.constants import AssetTypes\n",
    "import time, datetime, os\n",
    "\n",
    "# print the sdk version - you many want to share this in the issue you will report if parts of this notebook don't work\n",
    "!pip show azure-ai-ml"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setup authentication\n",
    "\n",
    "We are using `DefaultAzureCredential` to get access to workspace. When an access token is needed, it requests one using multiple identities(`EnvironmentCredential, ManagedIdentityCredential, SharedTokenCacheCredential, VisualStudioCodeCredential, AzureCliCredential, AzurePowerShellCredential`) in turn, stopping when one provides a token.\n",
    "Reference [here](https://docs.microsoft.com/en-us/python/api/azure-identity/azure.identity.defaultazurecredential?view=azure-python) for more information.\n",
    "\n",
    "`DefaultAzureCredential` should be capable of handling most Azure SDK authentication scenarios. \n",
    "Reference [here](https://docs.microsoft.com/en-us/python/api/azure-identity/azure.identity?view=azure-python) for all available credentials if it does not work for you.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    credential = DefaultAzureCredential()\n",
    "    # Check if given credential can get token successfully.\n",
    "    credential.get_token(\"https://management.azure.com/.default\")\n",
    "except Exception as ex:\n",
    "    # Fall back to InteractiveBrowserCredential in case DefaultAzureCredential not work\n",
    "    credential = InteractiveBrowserCredential()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Connect to a workspace and registry\n",
    "\n",
    "Most samples create one client to connect to the workspace. However, in this sample, you need two clients. First client, called `ml_client_workspace`, will be used to connect to a workspace and run jobs or deploy endpoints. Second client, called `ml_client_registry` will be used to connect to the registry to create components, environments and models.\n",
    "\n",
    "Replace the following:\n",
    "* `<SUBSCRIPTION_ID>`\n",
    "* `<RESOURCE_GROUP>`\n",
    "* `<AML_WORKSPACE_NAME>`\n",
    "* `<REGISTRY_NAME>`\n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml_client_workspace = MLClient(\n",
    "    credential=credential,\n",
    "    subscription_id=\"<SUBSCRIPTION_ID>\",\n",
    "    resource_group_name=\"<RESOURCE_GROUP>\",\n",
    "    workspace_name=\"<AML_WORKSPACE_NAME>\",\n",
    ")\n",
    "print(ml_client_workspace)\n",
    "\n",
    "ml_client_registry = MLClient(credential=credential, registry_name=\"<REGISTRY_NAME>\")\n",
    "print(ml_client_registry)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create a version number and setup root directory \n",
    "Make sure that you set the version number to something unique if this notebook has been run before. You can use the timestamp to generate a unique version number, the sample code for which is commented out. This will prevent any name and version conflicts when creating assets.\n",
    "\n",
    "Set the root directory in which the YAML definitions of the components, environments, etc. are present."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "import sys\n",
    "\n",
    "# version = str(123456)\n",
    "version = str(int(time.time()))\n",
    "print(\"version: \", version)\n",
    "\n",
    "parent_dir = os.path.abspath(\n",
    "    os.path.join(\n",
    "        sys.path[0],\n",
    "        \"../../../../cli/jobs/pipelines-with-components/nyc_taxi_data_regression\",\n",
    "    )\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create environment in registry\n",
    "\n",
    "**Note:** In this step, we are creating a environment in registry and will use that in a job later. You can also create an environment in workspace or use an existing environment in registry or workspace to use in your job.  \n",
    "\n",
    "You will use a docker file to create the environment. The docker file has base python image and few python dependencies required to run Scikit Learn training jobs. This notebook: [../environment/environment.ipynb](../environment/environment.ipynb)has more samples for environment create.\n",
    "\n",
    "Note that we use the `ml_client_registry` client because we plan to create the environment in registry. The syntax for creating environment in a workspace or registry are identical. You just use a client that is specific to the target - workspace or registry."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env_docker_context = Environment(\n",
    "    build=BuildContext(path=os.path.join(parent_dir, \"env_train\")),\n",
    "    name=\"SKLearnEnv\",\n",
    "    version=version,\n",
    "    description=\"Scikit Learn environment\",\n",
    ")\n",
    "ml_client_registry.environments.create_or_update(env_docker_context)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get environment from registry\n",
    "\n",
    "Get the environment using the `ml_client_registry` client. The syntax for getting environment in a workspace or registry are identical. You just use a client that is specific to the target - workspace or registry.\n",
    "\n",
    "You will use this environment in the next step to create a component in the registry."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env_from_registry = ml_client_registry.environments.get(\n",
    "    name=\"SKLearnEnv\", version=version\n",
    ")\n",
    "print(env_from_registry)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create component in registry\n",
    "\n",
    "**Note:** In this step, we are creating a component in registry and will use that in a job later. You can also create an component in workspace or use an existing component in registry or workspace to use in your job.  \n",
    "\n",
    "\n",
    "You will use the [`train.yml`](../../../cli/jobs/pipelines-with-components/nyc_taxi_data_regression/train.yml) component YAML defined in `cli/jobs/pipelines-with-components/nyc_taxi_data_regression` for this. This component runs a Scikit Learn training python script. The `train.yml` refers to the AzureML curated environment for the Scikit Learn framework: `AzureML-sklearn-0.24-ubuntu18`, but you will over ride this to use the Scikit Learn environment you created in the previous step.\n",
    "\n",
    "A similar sample notebook shows how to create these components in workspaces instead of registry, in which case you can use those components only in the specific workspace: https://github.com/Azure/azureml-examples/blob/main/sdk/jobs/pipelines/1e_pipeline_with_registered_components/pipeline_with_registered_components.ipynb\n",
    "\n",
    "\n",
    "Use the `ml_client_registry` client to create the component in the registry. The syntax for creating component in a workspace or registry are identical. You just use a client that is specific to the target - workspace or registry."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load component definition from YAML\n",
    "print(parent_dir)\n",
    "train_model = load_component(source=os.path.join(parent_dir, \"train.yml\"))\n",
    "# print the component as yaml\n",
    "print(train_model)\n",
    "\n",
    "# change environment reference to the environment created in registry\n",
    "train_model.environment = env_from_registry\n",
    "\n",
    "# changing the version number is optional, but useful if a component with same name and version already exist in registry\n",
    "train_model.version = version\n",
    "\n",
    "print(train_model)\n",
    "ml_client_registry.components.create_or_update(train_model)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get component from Registry\n",
    "\n",
    "Get the component using the `ml_client_registry` client. The syntax for getting component from a workspace or registry are identical. You just use a client that is specific to the target - workspace or registry.\n",
    "\n",
    "You will use this component in the next step to run a pipeline job to train a model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_component_from_registry = ml_client_registry.components.get(\n",
    "    name=\"train_linear_regression_model\", version=version\n",
    ")\n",
    "print(train_component_from_registry)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Create data asset in registry\n",
    "\n",
    "The data asset created in this step is used later in this article when submitting a training job. We will use transformed NYC taxi data which is required as in input to the pipelines job we will be running later. The data is avalble in `data_transformed` folder.\n",
    "\n",
    "Use the `ml_client_registry` client to create the data in the registry. The syntax for creating data in a workspace or registry are identical. You just use a client that is specific to the target - workspace or registry."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "my_path = parent_dir + \"/data_transformed/\"\n",
    "my_data = Data(\n",
    "    path=my_path,\n",
    "    type=AssetTypes.URI_FOLDER,\n",
    "    description=\"Transformed NYC Taxi data created from local folder.\",\n",
    "    name=\"transformed-nyc-taxt-data\",\n",
    "    version=version,\n",
    ")\n",
    "ml_client_registry.data.create_or_update(my_data)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Get data from registry\n",
    "\n",
    "Get the data using the ml_client_registry client. The syntax for getting data from a workspace or registry are identical. You just use a client that is specific to the target - workspace or registry.\n",
    "\n",
    "You will use this data in the next step to run a pipeline job to train a model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get the data asset\n",
    "data_asset_from_registry = ml_client_registry.data.get(\n",
    "    name=\"transformed-nyc-taxt-data\", version=version\n",
    ")\n",
    "print(data_asset_from_registry)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create a pipeline job using data from registry\n",
    "\n",
    "Review this page to learn how to use pipelines and components: https://github.com/Azure/azureml-examples/tree/main/sdk/python/jobs/pipelines. \n",
    "\n",
    "You will create a pipeline job that uses the training component created in the previous step using the Python DSL for pipelines. \n",
    "\n",
    "Make sure your workspace has a compute with the name `cpu-cluster` or update the compute name here: `pipeline_job.settings.default_compute = `"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@pipeline()\n",
    "def pipeline_with_registered_components(training_data):\n",
    "    train_job = train_component_from_registry(\n",
    "        training_data=training_data,\n",
    "    )\n",
    "\n",
    "\n",
    "pipeline_job = pipeline_with_registered_components(\n",
    "    training_data=Input(type=\"uri_folder\", path=data_asset_from_registry.id),\n",
    ")\n",
    "pipeline_job.settings.default_compute = \"cpu-cluster\"\n",
    "print(pipeline_job)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run pipeline job using a data from registry\n",
    "\n",
    "Submit pipeline job and wait for it to complete. Notice that you are using the workspace client: `ml_client_workspace` to run the pipeline job. This job is running a component and data that is not available in your workspace but is coming from a registry. This way, you can run this job in any workspace you have access to. This is useful when you want to develop a pipeline in th `dev` workspace with some sample data and run the pipeline in the `prod` workspace with actual data. This is also helpful if you want to share the components and data you develop with other teams in your organization who may be using a different workspace. \n",
    "To summarize, you can submit this job to different workspaces such as `dev`, `test` or `prod` by creating different ML clients for each of those workspaces."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pipeline_job = ml_client_workspace.jobs.create_or_update(\n",
    "    pipeline_job,\n",
    "    experiment_name=\"sdk_job_data_from_registry\",\n",
    "    skip_validation=True,\n",
    ")\n",
    "ml_client_workspace.jobs.stream(pipeline_job.name)\n",
    "pipeline_job = ml_client_workspace.jobs.get(pipeline_job.name)\n",
    "pipeline_job"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Share data from workspce to registry\n",
    "\n",
    "You can also share an existing data asset from workspace to registry. We will start with creating a data asset in workspace. For this we will use the same data from a previous step in this notebook but instead of creating it in registry, we will create it in workspace.\n",
    "\n",
    "Note that the only diffeence is that we are using `ml_clent_workspace` instead of `ml_client_registry`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "my_path = parent_dir + \"/data_transformed/\"\n",
    "my_data = Data(\n",
    "    path=my_path,\n",
    "    type=AssetTypes.URI_FOLDER,\n",
    "    description=\"Transformed NYC Taxi data created from local folder.\",\n",
    "    name=\"transformed-nyc-taxt-data-ws\",\n",
    "    version=version,\n",
    ")\n",
    "ml_client_workspace.data.create_or_update(my_data)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this step, we will share the data, created in workspace, in the previous step to the registry.\n",
    "\n",
    "Note that we are passing `share_with_name` and `share_with_version` parameter in `share` function. These parameters are optional and if you do not pass these data will be shared with same name and version as in workspace."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sharing data from workspace to registry\n",
    "ml_client_workspace.data.share(\n",
    "    name=\"transformed-nyc-taxt-data-ws\",\n",
    "    version=version,\n",
    "    registry_name=\"<REGISTRY_NAME>\",\n",
    "    share_with_name=\"transformed-nyc-taxt-data-shared-from-ws\",\n",
    "    share_with_version=version,\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10 - SDK V2",
   "language": "python",
   "name": "python310-sdkv2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}