prefect-client 2.14.10__py3-none-any.whl → 2.14.11__py3-none-any.whl

prefect/infrastructure/kubernetes.py

@@ -3,6 +3,7 @@ import enum
 import json
 import math
 import os
+import shlex
 import time
 from contextlib import contextmanager
 from typing import TYPE_CHECKING, Any, Dict, Generator, List, Optional, Tuple, Union
@@ -359,6 +360,69 @@ class KubernetesJob(Infrastructure):
     def preview(self):
         return yaml.dump(self.build_job())
 
+    def get_corresponding_worker_type(self):
+        return "kubernetes"
+
+    async def generate_work_pool_base_job_template(self):
+        from prefect.workers.utilities import (
+            get_default_base_job_template_for_infrastructure_type,
+        )
+
+        base_job_template = await get_default_base_job_template_for_infrastructure_type(
+            self.get_corresponding_worker_type()
+        )
+        assert (
+            base_job_template is not None
+        ), "Failed to retrieve default base job template."
+        for key, value in self.dict(exclude_unset=True, exclude_defaults=True).items():
+            if key == "command":
+                base_job_template["variables"]["properties"]["command"]["default"] = (
+                    shlex.join(value)
+                )
+            elif key in [
+                "type",
+                "block_type_slug",
+                "_block_document_id",
+                "_block_document_name",
+                "_is_anonymous",
+                "job",
+                "customizations",
+            ]:
+                continue
+            elif key == "image_pull_policy":
+                base_job_template["variables"]["properties"]["image_pull_policy"][
+                    "default"
+                ] = value.value
+            elif key == "cluster_config":
+                base_job_template["variables"]["properties"]["cluster_config"][
+                    "default"
+                ] = {
+                    "$ref": {
+                        "block_document_id": str(self.cluster_config._block_document_id)
+                    }
+                }
+            elif key in base_job_template["variables"]["properties"]:
+                base_job_template["variables"]["properties"][key]["default"] = value
+            else:
+                self.logger.warning(
+                    f"Variable {key!r} is not supported by Kubernetes work pools."
+                    " Skipping."
+                )
+
+        custom_job_manifest = self.dict(exclude_unset=True, exclude_defaults=True).get(
+            "job"
+        )
+        if custom_job_manifest:
+            job_manifest = self.build_job()
+        else:
+            job_manifest = copy.deepcopy(
+                base_job_template["job_configuration"]["job_manifest"]
+            )
+            job_manifest = self.customizations.apply(job_manifest)
+        base_job_template["job_configuration"]["job_manifest"] = job_manifest
+
+        return base_job_template
+
     def build_job(self) -> KubernetesManifest:
         """Builds the Kubernetes Job Manifest"""
         job_manifest = copy.copy(self.job)

prefect/infrastructure/provisioners/cloud_run.py

@@ -14,13 +14,18 @@ from rich.panel import Panel
 from rich.pretty import Pretty
 from rich.progress import Progress, SpinnerColumn, TextColumn
 from rich.prompt import Confirm
+from rich.syntax import Syntax
 
-from prefect.cli._prompts import prompt_select_from_table
+from prefect.cli._prompts import prompt, prompt_select_from_table
 from prefect.client.orchestration import PrefectClient, ServerType
 from prefect.client.schemas.actions import BlockDocumentCreate
 from prefect.client.utilities import inject_client
 from prefect.exceptions import ObjectAlreadyExists
-from prefect.settings import PREFECT_DEBUG_MODE
+from prefect.settings import (
+    PREFECT_DEBUG_MODE,
+    PREFECT_DEFAULT_DOCKER_BUILD_NAMESPACE,
+    update_current_profile,
+)
 
 
 class CloudRunPushProvisioner:
@@ -28,6 +33,9 @@ class CloudRunPushProvisioner:
         self._console = Console()
         self._project = None
         self._region = None
+        self._service_account_name = "prefect-cloud-run"
+        self._credentials_block_name = None
+        self._image_repository_name = "prefect-images"
 
     @property
     def console(self):
@@ -107,7 +115,9 @@ class CloudRunPushProvisioner:
 
     async def _enable_cloud_run_api(self):
         try:
-            await self._run_command("gcloud services enable run.googleapis.com")
+            await self._run_command(
+                f"gcloud services enable run.googleapis.com --project={self._project}"
+            )
 
         except subprocess.CalledProcessError as e:
             raise RuntimeError(
@@ -118,7 +128,7 @@ class CloudRunPushProvisioner:
     async def _create_service_account(self):
         try:
             await self._run_command(
-                "gcloud iam service-accounts create prefect-cloud-run"
+                f"gcloud iam service-accounts create {self._service_account_name}"
                 ' --display-name "Prefect Cloud Run Service Account"'
             )
         except subprocess.CalledProcessError as e:
@@ -134,27 +144,29 @@ class CloudRunPushProvisioner:
             try:
                 await self._run_command(
                     "gcloud iam service-accounts keys create"
-                    f" {tmpdir}/prefect-cloud-run-key.json"
-                    f" --iam-account=prefect-cloud-run@{self._project}.iam.gserviceaccount.com"
+                    f" {tmpdir}/{self._service_account_name}-key.json"
+                    f" --iam-account={self._service_account_name}@{self._project}.iam.gserviceaccount.com"
                 )
             except subprocess.CalledProcessError as e:
                 raise RuntimeError(
                     "Error creating service account key. Please ensure you have the"
                     " necessary permissions."
                 ) from e
-            key = json.loads((Path(tmpdir) / "prefect-cloud-run-key.json").read_text())
+            key = json.loads(
+                (Path(tmpdir) / f"{self._service_account_name}-key.json").read_text()
+            )
         return key
 
     async def _assign_roles(self):
         try:
             await self._run_command(
                 "gcloud projects add-iam-policy-binding"
-                f' {self._project} --member="serviceAccount:prefect-cloud-run@{self._project}.iam.gserviceaccount.com"'
+                f' {self._project} --member="serviceAccount:{self._service_account_name}@{self._project}.iam.gserviceaccount.com"'
                 ' --role="roles/iam.serviceAccountUser"'
             )
             await self._run_command(
                 "gcloud projects add-iam-policy-binding"
-                f' {self._project} --member="serviceAccount:prefect-cloud-run@{self._project}.iam.gserviceaccount.com"'
+                f' {self._project} --member="serviceAccount:{self._service_account_name}@{self._project}.iam.gserviceaccount.com"'
                 ' --role="roles/run.developer"'
             )
         except subprocess.CalledProcessError as e:
@@ -163,8 +175,47 @@ class CloudRunPushProvisioner:
                 " necessary permissions."
             ) from e
 
+    async def _enable_artifact_registry_api(self):
+        try:
+            await self._run_command(
+                "gcloud services enable artifactregistry.googleapis.com"
+                f" --project={self._project}"
+            )
+        except subprocess.CalledProcessError as e:
+            raise RuntimeError(
+                "Error enabling Artifact Registry API. Please ensure you have the"
+                " necessary permissions."
+            ) from e
+
+    async def _create_artifact_registry_repository(self, repository_name: str):
+        try:
+            await self._run_command(
+                "gcloud artifacts repositories create"
+                f" {repository_name} --repository-format=docker"
+                f" --location={self._region} --project={self._project}"
+            )
+        except subprocess.CalledProcessError as e:
+            if "already exists" not in e.output.decode("utf-8"):
+                return
+            raise RuntimeError(
+                "Error creating Artifact Registry repository. Please ensure you have"
+                " the necessary permissions."
+            ) from e
+
+    async def _login_to_artifact_registry(self):
+        try:
+            await self._run_command(
+                f"gcloud auth configure-docker {self._region}-docker.pkg.dev"
+                f" --project={self._project}"
+            )
+        except subprocess.CalledProcessError as e:
+            raise RuntimeError(
+                "Error logging into Artifact Registry. Please ensure you have the"
+                " necessary permissions."
+            ) from e
+
     async def _create_gcp_credentials_block(
-        self, work_pool_name: str, key: dict, client: PrefectClient
+        self, block_document_name: str, key: dict, client: PrefectClient
     ) -> UUID:
         credentials_block_type = await client.read_block_type_by_slug("gcp-credentials")
 
@@ -177,7 +228,7 @@ class CloudRunPushProvisioner:
         try:
             block_doc = await client.create_block_document(
                 block_document=BlockDocumentCreate(
-                    name=f"{work_pool_name}-push-pool-credentials",
+                    name=block_document_name,
                     data={"service_account_info": key},
                     block_type_id=credentials_block_type.id,
                     block_schema_id=credentials_block_schema.id,
@@ -186,24 +237,13 @@ class CloudRunPushProvisioner:
             return block_doc.id
         except ObjectAlreadyExists:
             block_doc = await client.read_block_document_by_name(
-                name=f"{work_pool_name}-push-pool-credentials",
+                name=block_document_name,
                 block_type_slug="gcp-credentials",
             )
             return block_doc.id
 
-    @inject_client
-    async def provision(
-        self,
-        work_pool_name: str,
-        base_job_template: dict,
-        client: Optional[PrefectClient] = None,
-    ) -> Dict[str, Any]:
-        assert client, "Client injection failed"
-        await self._verify_gcloud_ready()
-        self._project = await self._get_project()
-        self._region = await self._get_default_region()
-
-        table = Panel(
+    async def _create_provision_table(self, work_pool_name: str, client: PrefectClient):
+        return Panel(
             dedent(
                 f"""\
                     Provisioning infrastructure for your work pool [blue]{work_pool_name}[/] will require:
@@ -211,32 +251,120 @@ class CloudRunPushProvisioner:
                         Updates in GCP project [blue]{self._project}[/] in region [blue]{self._region}[/]
 
                             - Activate the Cloud Run API for your project
-                            - Create a service account for managing Cloud Run jobs: [blue]prefect-cloud-run[/]
+                            - Activate the Artifact Registry API for your project
+                            - Create an Artifact Registry repository named [blue]{self._image_repository_name}[/]
+                            - Create a service account for managing Cloud Run jobs: [blue]{self._service_account_name}[/]
                                 - Service account will be granted the following roles:
                                     - Service Account User
                                     - Cloud Run Developer
-                            - Create a key for service account [blue]prefect-cloud-run[/]
+                            - Create a key for service account: [blue]{self._service_account_name}[/]
 
                         Updates in Prefect {"workspace" if client.server_type == ServerType.CLOUD else "server"}
 
-                            - Create GCP credentials block [blue]{work_pool_name}-push-pool-credentials[/] to store the service account key
-                    """
+                            - Create GCP credentials block to store the service account key: [blue]{self._credentials_block_name}[/]
+                """
             ),
             expand=False,
         )
+
+    async def _customize_resource_names(
+        self, work_pool_name: str, client: PrefectClient
+    ) -> bool:
+        self._service_account_name = prompt(
+            "Please enter a name for the service account",
+            default=self._service_account_name,
+        )
+        self._credentials_block_name = prompt(
+            "Please enter a name for the GCP credentials block",
+            default=self._credentials_block_name,
+        )
+        self._image_repository_name = prompt(
+            "Please enter a name for the Artifact Registry repository",
+            default=self._image_repository_name,
+        )
+        table = await self._create_provision_table(work_pool_name, client)
+        self._console.print(table)
+
+        return Confirm.ask(
+            "Proceed with infrastructure provisioning?", console=self._console
+        )
+
+    @inject_client
+    async def provision(
+        self,
+        work_pool_name: str,
+        base_job_template: dict,
+        client: Optional[PrefectClient] = None,
+    ) -> Dict[str, Any]:
+        assert client, "Client injection failed"
+        await self._verify_gcloud_ready()
+        self._project = await self._get_project()
+        self._region = await self._get_default_region()
+        self._credentials_block_name = f"{work_pool_name}-push-pool-credentials"
+
+        table = await self._create_provision_table(work_pool_name, client)
         self._console.print(table)
         if self._console.is_interactive:
-            if not Confirm.ask(
-                "Proceed with infrastructure provisioning?", console=self._console
+            chosen_option = prompt_select_from_table(
+                self._console,
+                "Proceed with infrastructure provisioning with default resource names?",
+                [
+                    {"header": "Options:", "key": "option"},
+                ],
+                [
+                    {
+                        "option": (
+                            "Yes, proceed with infrastructure provisioning with default"
+                            " resource names"
+                        )
+                    },
+                    {"option": "Customize resource names"},
+                    {"option": "Do not proceed with infrastructure provisioning"},
+                ],
+            )
+            if chosen_option["option"] == "Customize resource names":
+                if not await self._customize_resource_names(work_pool_name, client):
+                    return base_job_template
+
+            elif (
+                chosen_option["option"]
+                == "Do not proceed with infrastructure provisioning"
             ):
                 return base_job_template
+            elif (
+                chosen_option["option"]
+                != "Yes, proceed with infrastructure provisioning with default"
+                " resource names"
+            ):
+                # basically, we should never hit this. i'm concerned that we might change
+                # the options in the future and forget to update this check
+                raise ValueError(f"Invalid option selected: {chosen_option['option']}")
 
         with Progress(console=self._console) as progress:
-            task = progress.add_task("Provisioning Infrastructure", total=5)
+            task = progress.add_task("Provisioning Infrastructure", total=9)
             progress.console.print("Activating Cloud Run API")
             await self._enable_cloud_run_api()
             progress.advance(task)
 
+            progress.console.print("Activating Artifact Registry API")
+            await self._enable_artifact_registry_api()
+            progress.advance(task)
+
+            progress.console.print("Creating Artifact Registry repository")
+            await self._create_artifact_registry_repository(self._image_repository_name)
+            progress.advance(task)
+
+            progress.console.print("Configuring authentication to Artifact Registry")
+            await self._login_to_artifact_registry()
+            progress.advance(task)
+
+            progress.console.print("Setting default Docker build namespace")
+            default_docker_build_namespace = f"{self._region}-docker.pkg.dev/{self._project}/{self._image_repository_name}"
+            update_current_profile(
+                {PREFECT_DEFAULT_DOCKER_BUILD_NAMESPACE: default_docker_build_namespace}
+            )
+            progress.advance(task)
+
             progress.console.print("Creating service account")
             await self._create_service_account()
             progress.advance(task)
@@ -251,7 +379,7 @@ class CloudRunPushProvisioner:
 
             progress.console.print("Creating GCP credentials block")
             block_doc_id = await self._create_gcp_credentials_block(
-                work_pool_name, key, client
+                self._credentials_block_name, key, client
             )
             base_job_template_copy = deepcopy(base_job_template)
             base_job_template_copy["variables"]["properties"]["credentials"][
@@ -259,6 +387,44 @@ class CloudRunPushProvisioner:
             ] = {"$ref": {"block_document_id": str(block_doc_id)}}
             progress.advance(task)
 
+            self._console.print(
+                dedent(
+                    f"""\
+                    Your default Docker build namespace has been set to [blue]{default_docker_build_namespace!r}[/].
+                    Use any image name to build and push to this registry by default:
+                    """
+                ),
+                Panel(
+                    Syntax(
+                        dedent(
+                            f"""\
+                        from prefect import flow
+                        from prefect.deployments import DeploymentImage
+
+
+                        @flow(log_prints=True)
+                        def my_flow(name: str = "world"):
+                            print(f"Hello {{name}}! I'm a flow running in Cloud Run!")
+
+
+                        if __name__ == "__main__":
+                            my_flow.deploy(
+                                name="my-deployment",
+                                work_pool_name="{work_pool_name}",
+                                image=DeploymentImage(
+                                    name="my-image:latest",
+                                    platform="linux/amd64",
+                                )
+                            )"""
+                        ),
+                        "python",
+                        background_color="default",
+                    ),
+                    title="example_deploy_script.py",
+                    expand=False,
+                ),
+            )
+
         self._console.print(
             (
                 f"Infrastructure successfully provisioned for '{work_pool_name}' work"