Skip to content

prefect_gcp.cloud_run

DEPRECATION WARNING:

This module is deprecated as of March 2024 and will not be available after September 2024. It has been replaced by the Cloud Run and Cloud Run V2 workers, which offer enhanced functionality and better performance.

For upgrade instructions, see https://docs.prefect.io/latest/guides/upgrade-guide-agents-to-workers/.

Integrations with Google Cloud Run Job.

Note this module is experimental. The intefaces within may change without notice.

Examples:

Run a job using Google Cloud Run Jobs:
```python
CloudRunJob(
    image="gcr.io/my-project/my-image",
    region="us-east1",
    credentials=my_gcp_credentials
).run()
```

Run a job that runs the command `echo hello world` using Google Cloud Run Jobs:
```python
CloudRunJob(
    image="gcr.io/my-project/my-image",
    region="us-east1",
    credentials=my_gcp_credentials
    command=["echo", "hello world"]
).run()
```

Classes

CloudRunJob

Bases: Infrastructure

Infrastructure block used to run GCP Cloud Run Jobs.

Project name information is provided by the Credentials object, and should always be correct as long as the Credentials object is for the correct project.

Note this block is experimental. The interface may change without notice.

Source code in prefect_gcp/cloud_run.py
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
@deprecated_class(
    start_date="Mar 2024",
    help=(
        "Use the Cloud Run or Cloud Run v2 worker instead."
        " Refer to the upgrade guide for more information:"
        " https://docs.prefect.io/latest/guides/upgrade-guide-agents-to-workers/."
    ),
)
class CloudRunJob(Infrastructure):
    """
    <span class="badge-api experimental"/>

    Infrastructure block used to run GCP Cloud Run Jobs.

    Project name information is provided by the Credentials object, and should always
    be correct as long as the Credentials object is for the correct project.

    Note this block is experimental. The interface may change without notice.
    """

    _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://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(
        "cloud-run-job", description="The slug for this task type."
    )
    image: str = Field(
        ...,
        title="Image Name",
        description=(
            "The image to use for a new Cloud Run Job. This value must "
            "refer to an image within either Google Container Registry "
            "or Google Artifact Registry, like `gcr.io/<project_name>/<repo>/`."
        ),
    )
    region: str = Field(..., description="The region where the Cloud Run Job resides.")
    credentials: GcpCredentials  # cannot be Field; else it shows as Json

    # Job settings
    cpu: Optional[int] = Field(
        default=None,
        title="CPU",
        description=(
            "The amount of compute allocated to the Cloud Run Job. "
            "The int must be valid based on the rules specified at "
            "https://cloud.google.com/run/docs/configuring/cpu#setting-jobs ."
        ),
    )
    memory: Optional[int] = Field(
        default=None,
        title="Memory",
        description="The amount of memory allocated to the Cloud Run Job.",
    )
    memory_unit: Optional[Literal["G", "Gi", "M", "Mi"]] = Field(
        default=None,
        title="Memory Units",
        description=(
            "The unit of memory. See "
            "https://cloud.google.com/run/docs/configuring/memory-limits#setting "
            "for additional details."
        ),
    )
    vpc_connector_name: Optional[str] = Field(
        default=None,
        title="VPC Connector Name",
        description="The name of the VPC connector to use for the Cloud Run Job.",
    )
    args: Optional[List[str]] = Field(
        default=None,
        description=(
            "Arguments to be passed to your Cloud Run Job's entrypoint command."
        ),
    )
    env: Dict[str, str] = Field(
        default_factory=dict,
        description="Environment variables to be passed to your Cloud Run Job.",
    )

    # Cleanup behavior
    keep_job: Optional[bool] = Field(
        default=False,
        title="Keep Job After Completion",
        description="Keep the completed Cloud Run Job on Google Cloud Platform.",
    )
    timeout: Optional[int] = Field(
        default=600,
        gt=0,
        le=3600,
        title="Job Timeout",
        description=(
            "The length of time that Prefect will wait for a Cloud Run Job to complete "
            "before raising an exception."
        ),
    )
    max_retries: Optional[int] = Field(
        default=3,
        ge=0,
        le=10,
        title="Max Retries",
        description=(
            "The maximum retries setting specifies the number of times a task is "
            "allowed to restart in case of failure before being failed permanently."
        ),
    )
    # For private use
    _job_name: str = None
    _execution: Optional[Execution] = None

    @property
    def job_name(self):
        """Create a unique and valid job name."""

        if self._job_name is None:
            # get `repo` from `gcr.io/<project_name>/repo/other`
            components = self.image.split("/")
            image_name = components[2]
            # only alphanumeric and '-' allowed for a job name
            modified_image_name = image_name.replace(":", "-").replace(".", "-")
            # make 50 char limit for final job name, which will be '<name>-<uuid>'
            if len(modified_image_name) > 17:
                modified_image_name = modified_image_name[:17]
            name = f"{modified_image_name}-{uuid4().hex}"
            self._job_name = name

        return self._job_name

    @property
    def memory_string(self):
        """Returns the string expected for memory resources argument."""
        if self.memory and self.memory_unit:
            return str(self.memory) + self.memory_unit
        return None

    @validator("image")
    def _remove_image_spaces(cls, value):
        """Deal with spaces in image names."""
        if value is not None:
            return value.strip()

    @root_validator
    def _check_valid_memory(cls, values):
        """Make sure memory conforms to expected values for API.
        See: https://cloud.google.com/run/docs/configuring/memory-limits#setting
        """  # noqa
        if (values.get("memory") is not None and values.get("memory_unit") is None) or (
            values.get("memory_unit") is not None and values.get("memory") is None
        ):
            raise ValueError(
                "A memory value and unit must both be supplied to specify a memory"
                " value other than the default memory value."
            )
        return values

    def get_corresponding_worker_type(self) -> str:
        """Return the corresponding worker type for this infrastructure block."""
        return "cloud-run"

    async def generate_work_pool_base_job_template(self) -> dict:
        """
        Generate a base job template for a cloud-run work pool with the same
        configuration as this block.

        Returns:
            - dict: a base job template for a cloud-run work pool
        """
        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 generate default base job template for Cloud Run worker."
        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",
                "memory_unit",
            ]:
                continue
            elif key == "credentials":
                if not self.credentials._block_document_id:
                    raise BlockNotSavedError(
                        "It looks like you are trying to use a block that"
                        " has not been saved. Please call `.save` on your block"
                        " before publishing it as a work pool."
                    )
                base_job_template["variables"]["properties"]["credentials"][
                    "default"
                ] = {
                    "$ref": {
                        "block_document_id": str(self.credentials._block_document_id)
                    }
                }
            elif key == "memory" and self.memory_string:
                base_job_template["variables"]["properties"]["memory"][
                    "default"
                ] = self.memory_string
            elif key == "cpu" and self.cpu is not None:
                base_job_template["variables"]["properties"]["cpu"][
                    "default"
                ] = f"{self.cpu * 1000}m"
            elif key == "args":
                # Not a default variable, but we can add it to the template
                base_job_template["variables"]["properties"]["args"] = {
                    "title": "Arguments",
                    "type": "string",
                    "description": "Arguments to be passed to your Cloud Run Job's entrypoint command.",  # noqa
                    "default": value,
                }
                base_job_template["job_configuration"]["job_body"]["spec"]["template"][
                    "spec"
                ]["template"]["spec"]["containers"][0]["args"] = "{{ args }}"
            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 Cloud Run work pools."
                    " Skipping."
                )

        return base_job_template

    def _create_job_error(self, exc):
        """Provides a nicer error for 404s when trying to create a Cloud Run Job."""
        # TODO consider lookup table instead of the if/else,
        # also check for documented errors
        if exc.status_code == 404:
            raise RuntimeError(
                f"Failed to find resources at {exc.uri}. Confirm that region"
                f" '{self.region}' is the correct region for your Cloud Run Job and"
                f" that {self.credentials.project} is the correct GCP project. If"
                f" your project ID is not correct, you are using a Credentials block"
                f" with permissions for the wrong project."
            ) from exc
        raise exc

    def _job_run_submission_error(self, exc):
        """Provides a nicer error for 404s when submitting job runs."""
        if exc.status_code == 404:
            pat1 = r"The requested URL [^ ]+ was not found on this server"
            # pat2 = (
            #     r"Resource '[^ ]+' of kind 'JOB' in region '[\w\-0-9]+' "
            #     r"in project '[\w\-0-9]+' does not exist"
            # )
            if re.findall(pat1, str(exc)):
                raise RuntimeError(
                    f"Failed to find resources at {exc.uri}. "
                    f"Confirm that region '{self.region}' is "
                    f"the correct region for your Cloud Run Job "
                    f"and that '{self.credentials.project}' is the "
                    f"correct GCP project. If your project ID is not "
                    f"correct, you are using a Credentials "
                    f"block with permissions for the wrong project."
                ) from exc
            else:
                raise exc

        raise exc

    def _cpu_as_k8s_quantity(self) -> str:
        """Return the CPU integer in the format expected by GCP Cloud Run Jobs API.
        See: https://kubernetes.io/docs/concepts/configuration/manage-resources-containers/
        See also: https://cloud.google.com/run/docs/configuring/cpu#setting-jobs
        """  # noqa
        return str(self.cpu * 1000) + "m"

    @sync_compatible
    async def run(self, task_status: Optional[TaskStatus] = None):
        """Run the configured job on a Google Cloud Run Job."""
        with self._get_client() as client:
            await run_sync_in_worker_thread(
                self._create_job_and_wait_for_registration, client
            )
            job_execution = await run_sync_in_worker_thread(
                self._begin_job_execution, client
            )

            if task_status:
                task_status.started(self.job_name)

            result = await run_sync_in_worker_thread(
                self._watch_job_execution_and_get_result,
                client,
                job_execution,
                5,
            )
            return result

    @sync_compatible
    async def kill(self, identifier: str, grace_seconds: int = 30) -> None:
        """
        Kill a task running Cloud Run.

        Args:
            identifier: The Cloud Run Job name. This should match a
                value yielded by CloudRunJob.run.
        """
        if grace_seconds != 30:
            self.logger.warning(
                f"Kill grace period of {grace_seconds}s requested, but GCP does not "
                "support dynamic grace period configuration. See here for more info: "
                "https://cloud.google.com/run/docs/reference/rest/v1/namespaces.jobs/delete"  # noqa
            )

        with self._get_client() as client:
            await run_sync_in_worker_thread(
                self._kill_job,
                client=client,
                namespace=self.credentials.project,
                job_name=identifier,
            )

    def _kill_job(self, client: Resource, namespace: str, job_name: str) -> None:
        """
        Thin wrapper around Job.delete, wrapping a try/except since
        Job is an independent class that doesn't have knowledge of
        CloudRunJob and its associated logic.
        """
        try:
            Job.delete(client=client, namespace=namespace, job_name=job_name)
        except Exception as exc:
            if "does not exist" in str(exc):
                raise InfrastructureNotFound(
                    f"Cannot stop Cloud Run Job; the job name {job_name!r} "
                    "could not be found."
                ) from exc
            raise

    def _create_job_and_wait_for_registration(self, client: Resource) -> None:
        """Create a new job wait for it to finish registering."""
        try:
            self.logger.info(f"Creating Cloud Run Job {self.job_name}")
            Job.create(
                client=client,
                namespace=self.credentials.project,
                body=self._jobs_body(),
            )
        except googleapiclient.errors.HttpError as exc:
            self._create_job_error(exc)

        try:
            self._wait_for_job_creation(client=client, timeout=self.timeout)
        except Exception:
            self.logger.exception(
                "Encountered an exception while waiting for job run creation"
            )
            if not self.keep_job:
                self.logger.info(
                    f"Deleting Cloud Run Job {self.job_name} from Google Cloud Run."
                )
                try:
                    Job.delete(
                        client=client,
                        namespace=self.credentials.project,
                        job_name=self.job_name,
                    )
                except Exception:
                    self.logger.exception(
                        "Received an unexpected exception while attempting to delete"
                        f" Cloud Run Job {self.job_name!r}"
                    )
            raise

    def _begin_job_execution(self, client: Resource) -> Execution:
        """Submit a job run for execution and return the execution object."""
        try:
            self.logger.info(
                f"Submitting Cloud Run Job {self.job_name!r} for execution."
            )
            submission = Job.run(
                client=client,
                namespace=self.credentials.project,
                job_name=self.job_name,
            )

            job_execution = Execution.get(
                client=client,
                namespace=submission["metadata"]["namespace"],
                execution_name=submission["metadata"]["name"],
            )

            command = (
                " ".join(self.command) if self.command else "default container command"
            )

            self.logger.info(
                f"Cloud Run Job {self.job_name!r}: Running command {command!r}"
            )
        except Exception as exc:
            self._job_run_submission_error(exc)

        return job_execution

    def _watch_job_execution_and_get_result(
        self, client: Resource, execution: Execution, poll_interval: int
    ) -> CloudRunJobResult:
        """Wait for execution to complete and then return result."""
        try:
            job_execution = self._watch_job_execution(
                client=client,
                job_execution=execution,
                timeout=self.timeout,
                poll_interval=poll_interval,
            )
        except Exception:
            self.logger.exception(
                "Received an unexpected exception while monitoring Cloud Run Job "
                f"{self.job_name!r}"
            )
            raise

        if job_execution.succeeded():
            status_code = 0
            self.logger.info(f"Job Run {self.job_name} completed successfully")
        else:
            status_code = 1
            error_msg = job_execution.condition_after_completion()["message"]
            self.logger.error(
                f"Job Run {self.job_name} did not complete successfully. {error_msg}"
            )

        self.logger.info(
            f"Job Run logs can be found on GCP at: {job_execution.log_uri}"
        )

        if not self.keep_job:
            self.logger.info(
                f"Deleting completed Cloud Run Job {self.job_name!r} from Google Cloud"
                " Run..."
            )
            try:
                Job.delete(
                    client=client,
                    namespace=self.credentials.project,
                    job_name=self.job_name,
                )
            except Exception:
                self.logger.exception(
                    "Received an unexpected exception while attempting to delete Cloud"
                    f" Run Job {self.job_name}"
                )

        return CloudRunJobResult(identifier=self.job_name, status_code=status_code)

    def _jobs_body(self) -> dict:
        """Create properly formatted body used for a Job CREATE request.
        See: https://cloud.google.com/run/docs/reference/rest/v1/namespaces.jobs
        """
        jobs_metadata = {"name": self.job_name}

        annotations = {
            # See: https://cloud.google.com/run/docs/troubleshooting#launch-stage-validation  # noqa
            "run.googleapis.com/launch-stage": "BETA",
        }
        # add vpc connector if specified
        if self.vpc_connector_name:
            annotations[
                "run.googleapis.com/vpc-access-connector"
            ] = self.vpc_connector_name

        # env and command here
        containers = [self._add_container_settings({"image": self.image})]

        # apply this timeout to each task
        timeout_seconds = str(self.timeout)

        body = {
            "apiVersion": "run.googleapis.com/v1",
            "kind": "Job",
            "metadata": jobs_metadata,
            "spec": {  # JobSpec
                "template": {  # ExecutionTemplateSpec
                    "metadata": {"annotations": annotations},
                    "spec": {  # ExecutionSpec
                        "template": {  # TaskTemplateSpec
                            "spec": {
                                "containers": containers,
                                "timeoutSeconds": timeout_seconds,
                                "maxRetries": self.max_retries,
                            }  # TaskSpec
                        }
                    },
                }
            },
        }
        return body

    def preview(self) -> str:
        """Generate a preview of the job definition that will be sent to GCP."""
        body = self._jobs_body()
        container_settings = body["spec"]["template"]["spec"]["template"]["spec"][
            "containers"
        ][0]["env"]
        body["spec"]["template"]["spec"]["template"]["spec"]["containers"][0]["env"] = [
            container_setting
            for container_setting in container_settings
            if container_setting["name"] != "PREFECT_API_KEY"
        ]
        return json.dumps(body, indent=2)

    def _watch_job_execution(
        self, client, job_execution: Execution, timeout: int, poll_interval: int = 5
    ):
        """
        Update job_execution status until it is no longer running or timeout is reached.
        """
        t0 = time.time()
        while job_execution.is_running():
            job_execution = Execution.get(
                client=client,
                namespace=job_execution.namespace,
                execution_name=job_execution.name,
            )

            elapsed_time = time.time() - t0
            if timeout is not None and elapsed_time > timeout:
                raise RuntimeError(
                    f"Timed out after {elapsed_time}s while waiting for Cloud Run Job "
                    "execution to complete. Your job may still be running on GCP."
                )

            time.sleep(poll_interval)

        return job_execution

    def _wait_for_job_creation(
        self, client: Resource, timeout: int, poll_interval: int = 5
    ):
        """Give created job time to register."""
        job = Job.get(
            client=client, namespace=self.credentials.project, job_name=self.job_name
        )

        t0 = time.time()
        while not job.is_ready():
            ready_condition = (
                job.ready_condition
                if job.ready_condition
                else "waiting for condition update"
            )
            self.logger.info(
                f"Job is not yet ready... Current condition: {ready_condition}"
            )
            job = Job.get(
                client=client,
                namespace=self.credentials.project,
                job_name=self.job_name,
            )

            elapsed_time = time.time() - t0
            if timeout is not None and elapsed_time > timeout:
                raise RuntimeError(
                    f"Timed out after {elapsed_time}s while waiting for Cloud Run Job "
                    "execution to complete. Your job may still be running on GCP."
                )

            time.sleep(poll_interval)

    def _get_client(self) -> Resource:
        """Get the base client needed for interacting with GCP APIs."""
        # region needed for 'v1' API
        api_endpoint = f"https://{self.region}-run.googleapis.com"
        gcp_creds = self.credentials.get_credentials_from_service_account()
        options = ClientOptions(api_endpoint=api_endpoint)

        return discovery.build(
            "run", "v1", client_options=options, credentials=gcp_creds
        ).namespaces()

    # CONTAINER SETTINGS
    def _add_container_settings(self, base_settings: Dict[str, Any]) -> Dict[str, Any]:
        """
        Add settings related to containers for Cloud Run Jobs to a dictionary.
        Includes environment variables, entrypoint command, entrypoint arguments,
        and cpu and memory limits.
        See: https://cloud.google.com/run/docs/reference/rest/v1/Container
        and https://cloud.google.com/run/docs/reference/rest/v1/Container#ResourceRequirements
        """  # noqa
        container_settings = base_settings.copy()
        container_settings.update(self._add_env())
        container_settings.update(self._add_resources())
        container_settings.update(self._add_command())
        container_settings.update(self._add_args())
        return container_settings

    def _add_args(self) -> dict:
        """Set the arguments that will be passed to the entrypoint for a Cloud Run Job.
        See: https://cloud.google.com/run/docs/reference/rest/v1/Container
        """  # noqa
        return {"args": self.args} if self.args else {}

    def _add_command(self) -> dict:
        """Set the command that a container will run for a Cloud Run Job.
        See: https://cloud.google.com/run/docs/reference/rest/v1/Container
        """  # noqa
        return {"command": self.command}

    def _add_resources(self) -> dict:
        """Set specified resources limits for a Cloud Run Job.
        See: https://cloud.google.com/run/docs/reference/rest/v1/Container#ResourceRequirements
        See also: https://kubernetes.io/docs/concepts/configuration/manage-resources-containers/
        """  # noqa
        resources = {"limits": {}, "requests": {}}

        if self.cpu is not None:
            cpu = self._cpu_as_k8s_quantity()
            resources["limits"]["cpu"] = cpu
            resources["requests"]["cpu"] = cpu
        if self.memory_string is not None:
            resources["limits"]["memory"] = self.memory_string
            resources["requests"]["memory"] = self.memory_string

        return {"resources": resources} if resources["requests"] else {}

    def _add_env(self) -> dict:
        """Add environment variables for a Cloud Run Job.

        Method `self._base_environment()` gets necessary Prefect environment variables
        from the config.

        See: https://cloud.google.com/run/docs/reference/rest/v1/Container#envvar for
        how environment variables are specified for Cloud Run Jobs.
        """  # noqa
        env = {**self._base_environment(), **self.env}
        cloud_run_env = [{"name": k, "value": v} for k, v in env.items()]
        return {"env": cloud_run_env}

Attributes

job_name property

Create a unique and valid job name.

memory_string property

Returns the string expected for memory resources argument.

Functions

generate_work_pool_base_job_template async

Generate a base job template for a cloud-run work pool with the same configuration as this block.

Returns:

Type Description
dict
  • dict: a base job template for a cloud-run work pool
Source code in prefect_gcp/cloud_run.py
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
async def generate_work_pool_base_job_template(self) -> dict:
    """
    Generate a base job template for a cloud-run work pool with the same
    configuration as this block.

    Returns:
        - dict: a base job template for a cloud-run work pool
    """
    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 generate default base job template for Cloud Run worker."
    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",
            "memory_unit",
        ]:
            continue
        elif key == "credentials":
            if not self.credentials._block_document_id:
                raise BlockNotSavedError(
                    "It looks like you are trying to use a block that"
                    " has not been saved. Please call `.save` on your block"
                    " before publishing it as a work pool."
                )
            base_job_template["variables"]["properties"]["credentials"][
                "default"
            ] = {
                "$ref": {
                    "block_document_id": str(self.credentials._block_document_id)
                }
            }
        elif key == "memory" and self.memory_string:
            base_job_template["variables"]["properties"]["memory"][
                "default"
            ] = self.memory_string
        elif key == "cpu" and self.cpu is not None:
            base_job_template["variables"]["properties"]["cpu"][
                "default"
            ] = f"{self.cpu * 1000}m"
        elif key == "args":
            # Not a default variable, but we can add it to the template
            base_job_template["variables"]["properties"]["args"] = {
                "title": "Arguments",
                "type": "string",
                "description": "Arguments to be passed to your Cloud Run Job's entrypoint command.",  # noqa
                "default": value,
            }
            base_job_template["job_configuration"]["job_body"]["spec"]["template"][
                "spec"
            ]["template"]["spec"]["containers"][0]["args"] = "{{ args }}"
        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 Cloud Run work pools."
                " Skipping."
            )

    return base_job_template
get_corresponding_worker_type

Return the corresponding worker type for this infrastructure block.

Source code in prefect_gcp/cloud_run.py
375
376
377
def get_corresponding_worker_type(self) -> str:
    """Return the corresponding worker type for this infrastructure block."""
    return "cloud-run"
kill async

Kill a task running Cloud Run.

Parameters:

Name Type Description Default
identifier str

The Cloud Run Job name. This should match a value yielded by CloudRunJob.run.

required
Source code in prefect_gcp/cloud_run.py
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
@sync_compatible
async def kill(self, identifier: str, grace_seconds: int = 30) -> None:
    """
    Kill a task running Cloud Run.

    Args:
        identifier: The Cloud Run Job name. This should match a
            value yielded by CloudRunJob.run.
    """
    if grace_seconds != 30:
        self.logger.warning(
            f"Kill grace period of {grace_seconds}s requested, but GCP does not "
            "support dynamic grace period configuration. See here for more info: "
            "https://cloud.google.com/run/docs/reference/rest/v1/namespaces.jobs/delete"  # noqa
        )

    with self._get_client() as client:
        await run_sync_in_worker_thread(
            self._kill_job,
            client=client,
            namespace=self.credentials.project,
            job_name=identifier,
        )
preview

Generate a preview of the job definition that will be sent to GCP.

Source code in prefect_gcp/cloud_run.py
715
716
717
718
719
720
721
722
723
724
725
726
def preview(self) -> str:
    """Generate a preview of the job definition that will be sent to GCP."""
    body = self._jobs_body()
    container_settings = body["spec"]["template"]["spec"]["template"]["spec"][
        "containers"
    ][0]["env"]
    body["spec"]["template"]["spec"]["template"]["spec"]["containers"][0]["env"] = [
        container_setting
        for container_setting in container_settings
        if container_setting["name"] != "PREFECT_API_KEY"
    ]
    return json.dumps(body, indent=2)
run async

Run the configured job on a Google Cloud Run Job.

Source code in prefect_gcp/cloud_run.py
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
@sync_compatible
async def run(self, task_status: Optional[TaskStatus] = None):
    """Run the configured job on a Google Cloud Run Job."""
    with self._get_client() as client:
        await run_sync_in_worker_thread(
            self._create_job_and_wait_for_registration, client
        )
        job_execution = await run_sync_in_worker_thread(
            self._begin_job_execution, client
        )

        if task_status:
            task_status.started(self.job_name)

        result = await run_sync_in_worker_thread(
            self._watch_job_execution_and_get_result,
            client,
            job_execution,
            5,
        )
        return result

CloudRunJobResult

Bases: InfrastructureResult

Result from a Cloud Run Job.

Source code in prefect_gcp/cloud_run.py
215
216
class CloudRunJobResult(InfrastructureResult):
    """Result from a Cloud Run Job."""

Execution

Bases: BaseModel

Utility class to call GCP executions API and interact with the returned objects.

Source code in prefect_gcp/cloud_run.py
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
class Execution(BaseModel):
    """
    Utility class to call GCP `executions` API and
    interact with the returned objects.
    """

    name: str
    namespace: str
    metadata: dict
    spec: dict
    status: dict
    log_uri: str

    def is_running(self) -> bool:
        """Returns True if Execution is not completed."""
        return self.status.get("completionTime") is None

    def condition_after_completion(self):
        """Returns Execution condition if Execution has completed."""
        for condition in self.status["conditions"]:
            if condition["type"] == "Completed":
                return condition

    def succeeded(self):
        """Whether or not the Execution completed is a successful state."""
        completed_condition = self.condition_after_completion()
        if completed_condition and completed_condition["status"] == "True":
            return True

        return False

    @classmethod
    def get(cls, client: Resource, namespace: str, execution_name: str):
        """
        Make a get request to the GCP executions API
        and return an Execution instance.
        """
        request = client.executions().get(
            name=f"namespaces/{namespace}/executions/{execution_name}"
        )
        response = request.execute()

        return cls(
            name=response["metadata"]["name"],
            namespace=response["metadata"]["namespace"],
            metadata=response["metadata"],
            spec=response["spec"],
            status=response["status"],
            log_uri=response["status"]["logUri"],
        )

Functions

condition_after_completion

Returns Execution condition if Execution has completed.

Source code in prefect_gcp/cloud_run.py
180
181
182
183
184
def condition_after_completion(self):
    """Returns Execution condition if Execution has completed."""
    for condition in self.status["conditions"]:
        if condition["type"] == "Completed":
            return condition
get classmethod

Make a get request to the GCP executions API and return an Execution instance.

Source code in prefect_gcp/cloud_run.py
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
@classmethod
def get(cls, client: Resource, namespace: str, execution_name: str):
    """
    Make a get request to the GCP executions API
    and return an Execution instance.
    """
    request = client.executions().get(
        name=f"namespaces/{namespace}/executions/{execution_name}"
    )
    response = request.execute()

    return cls(
        name=response["metadata"]["name"],
        namespace=response["metadata"]["namespace"],
        metadata=response["metadata"],
        spec=response["spec"],
        status=response["status"],
        log_uri=response["status"]["logUri"],
    )
is_running

Returns True if Execution is not completed.

Source code in prefect_gcp/cloud_run.py
176
177
178
def is_running(self) -> bool:
    """Returns True if Execution is not completed."""
    return self.status.get("completionTime") is None
succeeded

Whether or not the Execution completed is a successful state.

Source code in prefect_gcp/cloud_run.py
186
187
188
189
190
191
192
def succeeded(self):
    """Whether or not the Execution completed is a successful state."""
    completed_condition = self.condition_after_completion()
    if completed_condition and completed_condition["status"] == "True":
        return True

    return False

Job

Bases: BaseModel

Utility class to call GCP jobs API and interact with the returned objects.

Source code in prefect_gcp/cloud_run.py
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
class Job(BaseModel):
    """
    Utility class to call GCP `jobs` API and
    interact with the returned objects.
    """

    metadata: dict
    spec: dict
    status: dict
    name: str
    ready_condition: dict
    execution_status: dict

    def _is_missing_container(self):
        """
        Check if Job status is not ready because
        the specified container cannot be found.
        """
        if (
            self.ready_condition.get("status") == "False"
            and self.ready_condition.get("reason") == "ContainerMissing"
        ):
            return True
        return False

    def is_ready(self) -> bool:
        """Whether a job is finished registering and ready to be executed"""
        if self._is_missing_container():
            raise Exception(f"{self.ready_condition['message']}")
        return self.ready_condition.get("status") == "True"

    def has_execution_in_progress(self) -> bool:
        """See if job has a run in progress."""
        return (
            self.execution_status == {}
            or self.execution_status.get("completionTimestamp") is None
        )

    @staticmethod
    def _get_ready_condition(job: dict) -> dict:
        """Utility to access JSON field containing ready condition."""
        if job["status"].get("conditions"):
            for condition in job["status"]["conditions"]:
                if condition["type"] == "Ready":
                    return condition

        return {}

    @staticmethod
    def _get_execution_status(job: dict):
        """Utility to access JSON field containing execution status."""
        if job["status"].get("latestCreatedExecution"):
            return job["status"]["latestCreatedExecution"]

        return {}

    @classmethod
    def get(cls, client: Resource, namespace: str, job_name: str):
        """Make a get request to the GCP jobs API and return a Job instance."""
        request = client.jobs().get(name=f"namespaces/{namespace}/jobs/{job_name}")
        response = request.execute()

        return cls(
            metadata=response["metadata"],
            spec=response["spec"],
            status=response["status"],
            name=response["metadata"]["name"],
            ready_condition=cls._get_ready_condition(response),
            execution_status=cls._get_execution_status(response),
        )

    @staticmethod
    def create(client: Resource, namespace: str, body: dict):
        """Make a create request to the GCP jobs API."""
        request = client.jobs().create(parent=f"namespaces/{namespace}", body=body)
        response = request.execute()
        return response

    @staticmethod
    def delete(client: Resource, namespace: str, job_name: str):
        """Make a delete request to the GCP jobs API."""
        request = client.jobs().delete(name=f"namespaces/{namespace}/jobs/{job_name}")
        response = request.execute()
        return response

    @staticmethod
    def run(client: Resource, namespace: str, job_name: str):
        """Make a run request to the GCP jobs API."""
        request = client.jobs().run(name=f"namespaces/{namespace}/jobs/{job_name}")
        response = request.execute()
        return response

Functions

create staticmethod

Make a create request to the GCP jobs API.

Source code in prefect_gcp/cloud_run.py
141
142
143
144
145
146
@staticmethod
def create(client: Resource, namespace: str, body: dict):
    """Make a create request to the GCP jobs API."""
    request = client.jobs().create(parent=f"namespaces/{namespace}", body=body)
    response = request.execute()
    return response
delete staticmethod

Make a delete request to the GCP jobs API.

Source code in prefect_gcp/cloud_run.py
148
149
150
151
152
153
@staticmethod
def delete(client: Resource, namespace: str, job_name: str):
    """Make a delete request to the GCP jobs API."""
    request = client.jobs().delete(name=f"namespaces/{namespace}/jobs/{job_name}")
    response = request.execute()
    return response
get classmethod

Make a get request to the GCP jobs API and return a Job instance.

Source code in prefect_gcp/cloud_run.py
126
127
128
129
130
131
132
133
134
135
136
137
138
139
@classmethod
def get(cls, client: Resource, namespace: str, job_name: str):
    """Make a get request to the GCP jobs API and return a Job instance."""
    request = client.jobs().get(name=f"namespaces/{namespace}/jobs/{job_name}")
    response = request.execute()

    return cls(
        metadata=response["metadata"],
        spec=response["spec"],
        status=response["status"],
        name=response["metadata"]["name"],
        ready_condition=cls._get_ready_condition(response),
        execution_status=cls._get_execution_status(response),
    )
has_execution_in_progress

See if job has a run in progress.

Source code in prefect_gcp/cloud_run.py
101
102
103
104
105
106
def has_execution_in_progress(self) -> bool:
    """See if job has a run in progress."""
    return (
        self.execution_status == {}
        or self.execution_status.get("completionTimestamp") is None
    )
is_ready

Whether a job is finished registering and ready to be executed

Source code in prefect_gcp/cloud_run.py
95
96
97
98
99
def is_ready(self) -> bool:
    """Whether a job is finished registering and ready to be executed"""
    if self._is_missing_container():
        raise Exception(f"{self.ready_condition['message']}")
    return self.ready_condition.get("status") == "True"
run staticmethod

Make a run request to the GCP jobs API.

Source code in prefect_gcp/cloud_run.py
155
156
157
158
159
160
@staticmethod
def run(client: Resource, namespace: str, job_name: str):
    """Make a run request to the GCP jobs API."""
    request = client.jobs().run(name=f"namespaces/{namespace}/jobs/{job_name}")
    response = request.execute()
    return response