sdk/python/assets/assets-in-registry/share-models-components-environments.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Share components, environments and models 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-models-pipelines-across-workspaces-with-registries \n",
    "\n",
    "### Prerequisites\n",
    "Review the prerequisites section in the article: https://learn.microsoft.com/en-us/azure/machine-learning/how-to-share-models-pipelines-across-workspaces-with-registries?tabs=python#prerequisites. To summarize, in addition to the Python SDK, you need an AzureML registry and an AzureML workspace in a region that is supported by the workspace.\n",
    "\n",
    "\n",
    "### Scenarios\n",
    "\n",
    "There are two scenarios where you'd want to use the same set of models, components and environments in multiple workspaces:\n",
    "\n",
    "* Cross-workspace MLOps: You're training a model in a dev workspace and need to deploy it to test and prod workspaces. In this case you, want to have end-to-end lineage between endpoints to which the model is deployed in test or prod workspaces and the training job, metrics, code, data and environment that was used to train the model in the dev workspace.\n",
    "* Share and reuse models and pipelines across different teams: Sharing and reuse improve collaboration and productivity. In this scenario, you may want to publish a trained model and the associated components and environments used to train it to a central catalog. From there, colleagues from other teams can search and reuse the assets you shared in their own experiments.\n",
    "\n",
    "### Goals\n",
    "* Create an environment and component in the registry.\n",
    "* Use the component from registry to submit a model training job in a workspace.\n",
    "* Register the trained model in the registry.\n",
    "* Deploy the model from the registry to an online-endpoint in the workspace, then submit an inference request.\n"
   ]
  },
  {
   "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",
    "    ManagedOnlineEndpoint,\n",
    "    ManagedOnlineDeployment,\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"
   ]
  },
  {
   "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()"
   ]
  },
  {
   "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)"
   ]
  },
  {
   "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",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create environment in registry\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)"
   ]
  },
  {
   "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)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create component in registry\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)"
   ]
  },
  {
   "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)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create a pipeline job using component 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=parent_dir + \"/data_transformed/\"),\n",
    ")\n",
    "pipeline_job.settings.default_compute = \"cpu-cluster\"\n",
    "print(pipeline_job)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run pipeline job using a component 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 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 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_component_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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create model in registry\n",
    "\n",
    "You will now obtain the model trained by the pipeline job in the above step and create the model in the registry. For completeness, we are showing two options here.\n",
    "* First option shows to create a model in registry from job output without downloading it. This option is recommended when you want to track the lineage between the training job and the model. You will create a model directly from the job output (without downloading it) in your workspace and then copy the model from workspace to registry. \n",
    "* Second option shows how to create a model in registry from local files. In this case you will download the model from the job output. This option is helpful if you have an existing model from some external source and want to host it in the registry to be shared with many workspaces.\n",
    "\n",
    "Review this notebook to learn the different model types and how to create them in a workspace: [../../assets/model/model.ipynb](../model/model.ipynb). In the below example, you will work with a `mlflow_model` that will help you deploy this model for inference without writing any scoring scripts."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create model in workspace and copy it to registry\n",
    "\n",
    "Step a: Get the model path from job output. Note that you use the `ml_client_workspace` to get the model path from job output in a workspace."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "jobs = ml_client_workspace.jobs.list(parent_job_name=pipeline_job.name)\n",
    "for job in jobs:\n",
    "    if job.display_name == \"train_job\":\n",
    "        print(job.name)\n",
    "        model_path_from_job = (\n",
    "            \"azureml://jobs/{job_name}/outputs/artifacts/paths/model\".format(\n",
    "                job_name=job.name\n",
    "            )\n",
    "        )\n",
    "\n",
    "print(model_path_from_job)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create model in workspace and copy it to registry\n",
    "Step b: Create model in workspace from job output. Note that you use the `ml_client_workspace` to create the model in workspace."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mlflow_model = Model(\n",
    "    path=model_path_from_job,\n",
    "    type=AssetTypes.MLFLOW_MODEL,\n",
    "    name=\"nyc-taxi-model\",\n",
    "    version=version,\n",
    "    description=\"MLflow model created from job output\",\n",
    ")\n",
    "print(mlflow_model)\n",
    "ml_client_workspace.models.create_or_update(mlflow_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create model in workspace and copy it to registry\n",
    "\n",
    "Step c: Get the model from workspace, prepare to copy it to registry and create the model registry using the `model_ready_to_copy` object. Note that you both the clients here: First, `ml_client_workspace` client to get the model from workspace and prepare the `model_ready_to_copy` object. Second, `ml_client_registry` to create the model in registry using the `model_ready_to_copy` object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fetch the model from workspace\n",
    "model_in_workspace = ml_client_workspace.models.get(\n",
    "    name=\"nyc-taxi-model\", version=version\n",
    ")\n",
    "print(\"workspace model:\\n\\n\", model_in_workspace)\n",
    "# change the format such that the registry understands the model (when you print the model_ready_to_copy object, notice the asset id\n",
    "model_ready_to_copy = ml_client_workspace.models._prepare_to_copy(model_in_workspace)\n",
    "print(\"\\n\\nmodel ready to copy:\\n\\n\", model_ready_to_copy)\n",
    "# copy the model from registry to workspace\n",
    "ml_client_registry.models.create_or_update(model_ready_to_copy).wait()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### [OPTIONAL] Crete model from local files\n",
    "\n",
    "Step a: Download the model from job output to a local folder. Note that you the `ml_client_workspace` client to download the model from workspace.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "jobs = ml_client_workspace.jobs.list(parent_job_name=pipeline_job.name)\n",
    "for job in jobs:\n",
    "    if job.display_name == \"train_job\":\n",
    "        print(job.name)\n",
    "        ml_client_workspace.jobs.download(job.name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### [OPTIONAL] Crete model from local files \n",
    "Step b: Create a model in registry from files in a local folder. Note that you use the `ml_client_registry` client to create the model in registry. The syntax for creating model in a workspace or registry are identical. You just use a client that is specific to the target - workspace or registry.\n",
    "\n",
    "> **Warning:** If you have successfully created a model in registry in the previous steps, this section will fail as the model with the name and version will already exist. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this section is optional, will fail if this model name and version is already created in the registry in the previous steps\n",
    "mlflow_model = Model(\n",
    "    path=\"./artifacts/model/\",\n",
    "    type=AssetTypes.MLFLOW_MODEL,\n",
    "    name=\"nyc-taxi-model\",\n",
    "    version=str(int(version) + 1),\n",
    "    description=\"MLflow model created from local path\",\n",
    ")\n",
    "print(mlflow_model)\n",
    "ml_client_registry.models.create_or_update(mlflow_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Deploy model from registry to online endpoint in workspace\n",
    "\n",
    "You will deploy the model to an online endpoint and submit some sample inference requests in this section. Note that just like jobs, endpoints that host models are specific to a workspace. You can deploy the a model from a registry to many workspaces. This helps you develop a model in `dev` workspace, share it with a registry, and then deploy it to `test` and `prod` workspaces. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get model from registry\n",
    "\n",
    "Use the `ml_client_registry` client to get the model created in previous section from 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": [
    "mlflow_model_from_registry = ml_client_registry.models.get(\n",
    "    name=\"nyc-taxi-model\", version=version\n",
    ")\n",
    "print(mlflow_model_from_registry)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create an online endpoint \n",
    "\n",
    "Create an online endpoint to deploy the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "online_endpoint_name = \"endpoint-\" + version\n",
    "# create an online endpoint\n",
    "endpoint = ManagedOnlineEndpoint(\n",
    "    name=online_endpoint_name,\n",
    "    description=\"this is a sample online endpoint for mlflow model\",\n",
    "    auth_mode=\"key\",\n",
    ")\n",
    "ml_client_workspace.begin_create_or_update(endpoint).wait()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Deploy the model from registry to the online endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a demo deployment\n",
    "demo_deployment = ManagedOnlineDeployment(\n",
    "    name=\"demo\",\n",
    "    endpoint_name=online_endpoint_name,\n",
    "    model=mlflow_model_from_registry.id,\n",
    "    instance_type=\"Standard_F4s_v2\",\n",
    "    instance_count=1,\n",
    ")\n",
    "ml_client_workspace.online_deployments.begin_create_or_update(demo_deployment).wait()\n",
    "\n",
    "endpoint.traffic = {\"demo\": 100}\n",
    "ml_client_workspace.begin_create_or_update(endpoint)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test the deployment\n",
    "\n",
    "This section needs a sample request file `scoring-data.json` which is available in the root directory initialized in the beginning of this notebook: [../../../cli/jobs/pipelines-with-components/nyc_taxi_data_regression/scoring-data.json](../../../cli/jobs/pipelines-with-components/nyc_taxi_data_regression/scoring-data.json)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test the  deployment with some sample data\n",
    "ml_client_workspace.online_endpoints.invoke(\n",
    "    endpoint_name=online_endpoint_name,\n",
    "    deployment_name=\"demo\",\n",
    "    request_file=parent_dir + \"/scoring-data.json\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Clean up resources\n",
    "\n",
    "#### delete online endpoint "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"online_endpoint_name: {online_endpoint_name}\")\n",
    "ml_client_workspace.online_endpoints.begin_delete(name=online_endpoint_name).wait()"
   ]
  }
 ],
 "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"
  },
  "vscode": {
   "interpreter": {
    "hash": "cb451b1acfd48ceac5a931ea70033c1eb64868bd3ea755a5c8b8f7161e90ef7b"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}