feat: Vertex AI CustomJob worker (#211)

* initial worker commit * Rename the module so we can import/register it; fix command-line formatting * make image required, more comments * tests * docs * remove unneeded defaults * oops * fix test, and rename file * combine to new workers module, fix docs * fix test and rename file * add configurable job spec, schema validation * creds * pass in creds as config variable * remove unneeded checks * fix doc links * set defaults on the variable config * be better at python * logogogogo * all logos * dont get cute --------- Co-authored-by: Chris Guidry <chris.g@prefect.io>
PrefectHQ · Sep 21, 2023 · 6327600 · 6327600
1 parent dfa07ca
commit 6327600
Show file tree

Hide file tree

Showing 16 changed files with 826 additions and 20 deletions.
diff --git a/docs/cloud_run_worker.md b/docs/cloud_run_worker.md
@@ -0,0 +1 @@
+::: prefect_gcp.workers.cloud_run
diff --git a/docs/vertex_worker.md b/docs/vertex_worker.md
@@ -0,0 +1 @@
+::: prefect_gcp.workers.vertex
diff --git a/docs/worker.md b/docs/worker.md
diff --git a/mkdocs.yml b/mkdocs.yml
@@ -85,7 +85,9 @@ nav:
         - Cloud Run: cloud_run.md
         - AI Platform: aiplatform.md
         - Deployment Steps: deployments/steps.md
-        - Worker: worker.md
+        - Workers:
+          - Cloud Run: cloud_run_worker.md
+          - Vertex AI: vertex_worker.md
 
 extra:
     social:

diff --git a/prefect_gcp/__init__.py b/prefect_gcp/__init__.py
@@ -1,19 +1,23 @@
-from . import _version
+from prefect._internal.compatibility.deprecated import (
+    register_renamed_module,
+)
 
-from .bigquery import BigQueryWarehouse  # noqa
+from . import _version
 from .aiplatform import VertexAICustomTrainingJob  # noqa
-from .cloud_storage import GcsBucket  # noqa
+from .bigquery import BigQueryWarehouse  # noqa
 from .cloud_run import CloudRunJob  # noqa
-from .secret_manager import GcpSecret  # noqa
+from .cloud_storage import GcsBucket  # noqa
 from .credentials import GcpCredentials  # noqa
-from .worker import CloudRunWorker  # noqa
-from prefect._internal.compatibility.deprecated import (
-    register_renamed_module,
-)
+from .secret_manager import GcpSecret  # noqa
+from .workers.vertex import VertexAIWorker  # noqa
+from .workers.cloud_run import CloudRunWorker  # noqa
 
 register_renamed_module(
     "prefect_gcp.projects", "prefect_gcp.deployments", start_date="Jun 2023"
 )
+register_renamed_module(
+    "prefect_gcp.worker", "prefect_gcp.workers", start_date="Sep 2023"
+)
 
 
 __version__ = _version.get_versions()["version"]
diff --git a/prefect_gcp/aiplatform.py b/prefect_gcp/aiplatform.py
@@ -102,7 +102,7 @@ class VertexAICustomTrainingJob(Infrastructure):
 
     _block_type_name = "Vertex AI Custom Training Job"
     _block_type_slug = "vertex-ai-custom-training-job"
-    _logo_url = "https://images.ctfassets.net/gm98wzqotmnx/4CD4wwbiIKPkZDt4U3TEuW/c112fe85653da054b6d5334ef662bec4/gcp.png?h=250"  # noqa
+    _logo_url = "https://cdn.sanity.io/images/3ugk85nk/production/10424e311932e31c477ac2b9ef3d53cefbaad708-250x250.png"  # noqa
     _documentation_url = "https://prefecthq.github.io/prefect-gcp/aiplatform/#prefect_gcp.aiplatform.VertexAICustomTrainingJob"  # noqa: E501
 
     type: Literal["vertex-ai-custom-training-job"] = Field(

diff --git a/prefect_gcp/bigquery.py b/prefect_gcp/bigquery.py
@@ -552,7 +552,7 @@ class BigQueryWarehouse(DatabaseBlock):
     """  # noqa
 
     _block_type_name = "BigQuery Warehouse"
-    _logo_url = "https://images.ctfassets.net/gm98wzqotmnx/4CD4wwbiIKPkZDt4U3TEuW/c112fe85653da054b6d5334ef662bec4/gcp.png?h=250"  # noqa
+    _logo_url = "https://cdn.sanity.io/images/3ugk85nk/production/10424e311932e31c477ac2b9ef3d53cefbaad708-250x250.png"  # noqa
     _documentation_url = "https://prefecthq.github.io/prefect-gcp/bigquery/#prefect_gcp.bigquery.BigQueryWarehouse"  # noqa: E501
 
     gcp_credentials: GcpCredentials

diff --git a/prefect_gcp/cloud_run.py b/prefect_gcp/cloud_run.py
@@ -213,7 +213,7 @@ class CloudRunJob(Infrastructure):
     _block_type_slug = "cloud-run-job"
     _block_type_name = "GCP Cloud Run Job"
     _description = "Infrastructure block used to run GCP Cloud Run Jobs. Note this block is experimental. The interface may change without notice."  # noqa
-    _logo_url = "https://images.ctfassets.net/gm98wzqotmnx/4CD4wwbiIKPkZDt4U3TEuW/c112fe85653da054b6d5334ef662bec4/gcp.png?h=250"  # noqa
+    _logo_url = "https://cdn.sanity.io/images/3ugk85nk/production/10424e311932e31c477ac2b9ef3d53cefbaad708-250x250.png"  # noqa
     _documentation_url = "https://prefecthq.github.io/prefect-gcp/cloud_run/#prefect_gcp.cloud_run.CloudRunJob"  # noqa: E501
 
     type: Literal["cloud-run-job"] = Field(

diff --git a/prefect_gcp/cloud_storage.py b/prefect_gcp/cloud_storage.py
@@ -564,7 +564,7 @@ class GcsBucket(WritableDeploymentStorage, WritableFileSystem, ObjectStorageBloc
         ```
     """
 
-    _logo_url = "https://images.ctfassets.net/gm98wzqotmnx/4CD4wwbiIKPkZDt4U3TEuW/c112fe85653da054b6d5334ef662bec4/gcp.png?h=250"  # noqa
+    _logo_url = "https://cdn.sanity.io/images/3ugk85nk/production/10424e311932e31c477ac2b9ef3d53cefbaad708-250x250.png"  # noqa
     _block_type_name = "GCS Bucket"
     _documentation_url = "https://prefecthq.github.io/prefect-gcp/cloud_storage/#prefect_gcp.cloud_storage.GcsBucket"  # noqa: E501
 

diff --git a/prefect_gcp/credentials.py b/prefect_gcp/credentials.py
@@ -96,7 +96,7 @@ class GcpCredentials(CredentialsBlock):
         ```
     """  # noqa
 
-    _logo_url = "https://images.ctfassets.net/gm98wzqotmnx/4CD4wwbiIKPkZDt4U3TEuW/c112fe85653da054b6d5334ef662bec4/gcp.png?h=250"  # noqa
+    _logo_url = "https://cdn.sanity.io/images/3ugk85nk/production/10424e311932e31c477ac2b9ef3d53cefbaad708-250x250.png"  # noqa
     _block_type_name = "GCP Credentials"
     _documentation_url = "https://prefecthq.github.io/prefect-gcp/credentials/#prefect_gcp.credentials.GcpCredentials"  # noqa: E501
 

diff --git a/prefect_gcp/secret_manager.py b/prefect_gcp/secret_manager.py
@@ -300,7 +300,7 @@ class GcpSecret(SecretBlock):
         secret_version: Version number of the secret to use, or "latest".
     """
 
-    _logo_url = "https://images.ctfassets.net/gm98wzqotmnx/4CD4wwbiIKPkZDt4U3TEuW/c112fe85653da054b6d5334ef662bec4/gcp.png?h=250"  # noqa
+    _logo_url = "https://cdn.sanity.io/images/3ugk85nk/production/10424e311932e31c477ac2b9ef3d53cefbaad708-250x250.png"  # noqa
     _documentation_url = "https://prefecthq.github.io/prefect-gcp/secret_manager/#prefect_gcp.secret_manager.GcpSecret"  # noqa: E501
 
     gcp_credentials: GcpCredentials

diff --git a/prefect_gcp/workers/__init__.py b/prefect_gcp/workers/__init__.py
diff --git a/prefect_gcp/worker.py → prefect_gcp/workers/cloud_run.py b/prefect_gcp/worker.py → prefect_gcp/workers/cloud_run.py
@@ -519,8 +519,8 @@ class CloudRunWorker(BaseWorker):
         "a Google Cloud Platform account."
     )
     _display_name = "Google Cloud Run"
-    _documentation_url = "https://prefecthq.github.io/prefect-gcp/worker/"
-    _logo_url = "https://images.ctfassets.net/gm98wzqotmnx/4SpnOBvMYkHp6z939MDKP6/549a91bc1ce9afd4fb12c68db7b68106/social-icon-google-cloud-1200-630.png?h=250"  # noqa
+    _documentation_url = "https://prefecthq.github.io/prefect-gcp/cloud_run_worker/"
+    _logo_url = "https://cdn.sanity.io/images/3ugk85nk/production/10424e311932e31c477ac2b9ef3d53cefbaad708-250x250.png"  # noqa
 
     def _create_job_error(self, exc, configuration):
         """Provides a nicer error for 404s when trying to create a Cloud Run Job."""