Skip to content

Commit 87af64a

Browse files
bcreddy-gcpmelinath
bcreddy-gcp
and
melinath
authored
Colab schedules (#12826)
Co-authored-by: Stephen Lewis (Burrows) <[email protected]>
1 parent 4dc815e commit 87af64a

File tree

9 files changed

+1133
-0
lines changed

9 files changed

+1133
-0
lines changed

‎mmv1/products/colab/Schedule.yaml

+211
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,211 @@
1+
# Copyright 2025 Google Inc.
2+
# Licensed under the Apache License, Version 2.0 (the "License");
3+
# you may not use this file except in compliance with the License.
4+
# You may obtain a copy of the License at
5+
#
6+
# http://www.apache.org/licenses/LICENSE-2.0
7+
#
8+
# Unless required by applicable law or agreed to in writing, software
9+
# distributed under the License is distributed on an "AS IS" BASIS,
10+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+
# See the License for the specific language governing permissions and
12+
# limitations under the License.
13+
14+
---
15+
16+
name: 'Schedule'
17+
description: |
18+
'Colab Enterprise Notebook Execution Schedules.'
19+
references:
20+
guides:
21+
'Schedule a notebook run': 'https://cloud.google.com/colab/docs/schedule-notebook-run'
22+
api: 'https://cloud.google.com/vertex-ai/docs/reference/rest/v1/projects.locations.schedules'
23+
base_url: 'projects/{{project}}/locations/{{location}}/schedules'
24+
self_link: 'projects/{{project}}/locations/{{location}}/schedules/{{name}}'
25+
create_url: 'projects/{{project}}/locations/{{location}}/schedules'
26+
update_url: 'projects/{{project}}/locations/{{location}}/schedules/{{name}}'
27+
update_mask: true
28+
update_verb: 'PATCH'
29+
async:
30+
type: 'OpAsync'
31+
actions: ['delete']
32+
operation:
33+
full_url: 'https://{{location}}-aiplatform.googleapis.com/v1/{{op_id}}'
34+
custom_code:
35+
encoder: 'templates/terraform/encoders/colab_schedule.go.tmpl'
36+
post_create: 'templates/terraform/post_create/colab_schedule.go.tmpl'
37+
post_update: 'templates/terraform/post_update/colab_schedule.go.tmpl'
38+
constants: 'templates/terraform/constants/colab_schedule.go.tmpl'
39+
sweeper:
40+
sweepable_identifier_field: 'displayName'
41+
examples:
42+
- name: 'colab_schedule_basic'
43+
primary_resource_id: 'schedule'
44+
vars:
45+
display_name: 'basic-schedule'
46+
bucket: 'my_bucket'
47+
runtime_template_name: 'runtime-template'
48+
test_env_vars:
49+
project_id: 'PROJECT_NAME'
50+
location: 'REGION'
51+
service_account: 'SERVICE_ACCT'
52+
- name: 'colab_schedule_paused'
53+
primary_resource_id: 'schedule'
54+
vars:
55+
display_name: 'paused-schedule'
56+
bucket: 'my_bucket'
57+
runtime_template_name: 'runtime-template'
58+
test_env_vars:
59+
project_id: 'PROJECT_NAME'
60+
location: 'REGION'
61+
service_account: 'SERVICE_ACCT'
62+
ignore_read_extra:
63+
- desired_state
64+
- name: 'colab_schedule_full'
65+
primary_resource_id: 'schedule'
66+
bootstrap_iam:
67+
- member: "serviceAccount:service-{project_number}@gcp-sa-dataform.iam.gserviceaccount.com"
68+
role: "roles/cloudkms.cryptoKeyEncrypterDecrypter"
69+
vars:
70+
display_name: 'full-schedule'
71+
bucket: 'my_bucket'
72+
network_name: 'colab-test-default'
73+
runtime_template_name: 'runtime-template'
74+
secret: 'secret'
75+
dataform_repository: 'dataform-repository'
76+
start_time: '2014-10-02T15:01:23Z'
77+
end_time: '2014-10-10T15:01:23Z'
78+
test_env_vars:
79+
project_id: 'PROJECT_NAME'
80+
location: 'REGION'
81+
service_account: 'SERVICE_ACCT'
82+
test_vars_overrides:
83+
key_name: 'acctest.BootstrapKMSKeyInLocation(t, "us-central1").CryptoKey.Name'
84+
# Tests use a (relatively static) future date so that VCR cassettes can stay valid for a long time
85+
start_time: 'time.Date(time.Now().Year(), 12, 31, 0, 0, 0, 0, time.Now().Location()).AddDate(0, 0, 1).Format(time.RFC3339)'
86+
end_time: 'time.Date(time.Now().Year(), 12, 31, 0, 0, 0, 0, time.Now().Location()).AddDate(0, 0, 10).Format(time.RFC3339)'
87+
virtual_fields:
88+
- name: 'desired_state'
89+
description: |
90+
Desired state of the Colab Schedule. Set this field to `ACTIVE` to start/resume the schedule, and `PAUSED` to pause the schedule.
91+
type: String
92+
default_value: "ACTIVE"
93+
parameters:
94+
- name: 'location'
95+
type: String
96+
required: true
97+
url_param_only: true
98+
description: 'The location for the resource: https://cloud.google.com/colab/docs/locations'
99+
properties:
100+
- name: 'name'
101+
type: String
102+
description: 'The resource name of the Schedule'
103+
custom_flatten: 'templates/terraform/custom_flatten/name_from_self_link.tmpl'
104+
output: true
105+
- name: 'displayName'
106+
type: String
107+
description:
108+
Required. The display name of the Schedule.
109+
required: true
110+
- name: 'startTime'
111+
type: String
112+
description:
113+
The timestamp after which the first run can be scheduled. Defaults to the schedule creation time. Must be in the RFC 3339 (https://www.ietf.org/rfc/rfc3339.txt) format.
114+
default_from_api: true
115+
- name: 'endTime'
116+
type: String
117+
description:
118+
Timestamp after which no new runs can be scheduled. If specified, the schedule will be completed when either end_time is reached or when scheduled_run_count >= max_run_count. Must be in the RFC 3339 (https://www.ietf.org/rfc/rfc3339.txt) format.
119+
- name: 'maxRunCount'
120+
type: int
121+
description:
122+
Maximum run count of the schedule. If specified, The schedule will be completed when either startedRunCount >= maxRunCount or when endTime is reached. If not specified, new runs will keep getting scheduled until this Schedule is paused or deleted. Already scheduled runs will be allowed to complete. Unset if not specified.
123+
- name: 'cron'
124+
type: string
125+
description: |
126+
Cron schedule (https://en.wikipedia.org/wiki/Cron) to launch scheduled runs.
127+
required: true
128+
- name: 'maxConcurrentRunCount'
129+
type: int
130+
required: true
131+
description:
132+
Maximum number of runs that can be started concurrently for this Schedule. This is the limit for starting the scheduled requests and not the execution of the notebook execution jobs created by the requests.
133+
- name: 'allowQueueing'
134+
type: Boolean
135+
send_empty_value: true
136+
description:
137+
Whether new scheduled runs can be queued when max_concurrent_runs limit is reached. If set to true, new runs will be queued instead of skipped. Default to false.
138+
- name: 'state'
139+
type: String
140+
description: |
141+
Output only. The state of the schedule.
142+
output: true
143+
- name: 'createNotebookExecutionJobRequest'
144+
immutable: true
145+
type: NestedObject
146+
description: Request for google_colab_notebook_execution.
147+
required: true
148+
properties:
149+
- name: 'notebookExecutionJob'
150+
type: NestedObject
151+
required: true
152+
description:
153+
The NotebookExecutionJob to create.
154+
properties:
155+
- name: 'displayName'
156+
type: String
157+
description:
158+
Required. The display name of the Notebook Execution.
159+
required: true
160+
- name: 'dataformRepositorySource'
161+
type: NestedObject
162+
description: 'The Dataform Repository containing the input notebook.'
163+
exactly_one_of:
164+
- create_notebook_execution_job_request.0.notebook_execution_job.0.dataform_repository_source
165+
- create_notebook_execution_job_request.0.notebook_execution_job.0.gcs_notebook_source
166+
properties:
167+
- name: 'dataformRepositoryResourceName'
168+
type: String
169+
required: true
170+
description: 'The resource name of the Dataform Repository.'
171+
diff_suppress_func: 'tpgresource.CompareSelfLinkRelativePaths'
172+
- name: 'commitSha'
173+
type: String
174+
description: 'The commit SHA to read repository with. If unset, the file will be read at HEAD.'
175+
- name: 'gcsNotebookSource'
176+
type: NestedObject
177+
description: 'The Cloud Storage uri for the input notebook.'
178+
exactly_one_of:
179+
- create_notebook_execution_job_request.0.notebook_execution_job.0.dataform_repository_source
180+
- create_notebook_execution_job_request.0.notebook_execution_job.0.gcs_notebook_source
181+
properties:
182+
- name: 'uri'
183+
type: String
184+
required: true
185+
description: 'The Cloud Storage uri pointing to the ipynb file. Format: gs://bucket/notebook_file.ipynb'
186+
- name: 'generation'
187+
type: String
188+
description: 'The version of the Cloud Storage object to read. If unset, the current version of the object is read. See https://cloud.google.com/storage/docs/metadata#generation-number.'
189+
- name: executionTimeout
190+
type: String
191+
description: 'Max running time of the execution job in seconds (default 86400s / 24 hrs). A duration in seconds with up to nine fractional digits, ending with "s". Example: "3.5s".'
192+
- name: 'notebookRuntimeTemplateResourceName'
193+
type: String
194+
description: 'The NotebookRuntimeTemplate to source compute configuration from.'
195+
required: true
196+
- name: 'gcsOutputUri'
197+
required: true
198+
type: String
199+
description: 'The Cloud Storage location to upload the result to. Format:`gs://bucket-name`'
200+
- name: 'executionUser'
201+
type: String
202+
exactly_one_of:
203+
- create_notebook_execution_job_request.0.notebook_execution_job.0.execution_user
204+
- create_notebook_execution_job_request.0.notebook_execution_job.0.service_account
205+
description: 'The user email to run the execution as.'
206+
- name: 'serviceAccount'
207+
type: String
208+
exactly_one_of:
209+
- create_notebook_execution_job_request.0.notebook_execution_job.0.execution_user
210+
- create_notebook_execution_job_request.0.notebook_execution_job.0.service_account
211+
description: 'The service account to run the execution as.'

‎mmv1/templates/terraform/constants/colab_schedule.go.tmpl

+18
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
func modifyScheduleState(config *transport_tpg.Config, d *schema.ResourceData, project string, billingProject string, userAgent string, state string) (map[string]interface{}, error) {
2+
url, err := tpgresource.ReplaceVars(d, config, "{{"{{"}}ColabBasePath{{"}}"}}projects/{{"{{"}}project{{"}}"}}/locations/{{"{{"}}location{{"}}"}}/schedules/{{"{{"}}name{{"}}"}}:"+state)
3+
if err != nil {
4+
return nil, err
5+
}
6+
7+
res, err := transport_tpg.SendRequest(transport_tpg.SendRequestOptions{
8+
Config: config,
9+
Method: "POST",
10+
Project: billingProject,
11+
RawURL: url,
12+
UserAgent: userAgent,
13+
})
14+
if err != nil {
15+
return nil, fmt.Errorf("Unable to %q google_colab_schedule %q: %s", state, d.Id(), err)
16+
}
17+
return res, nil
18+
}

‎mmv1/templates/terraform/encoders/colab_schedule.go.tmpl

+23
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
config := meta.(*transport_tpg.Config)
2+
project, err := tpgresource.GetProject(d, config)
3+
if err != nil {
4+
return nil, err
5+
}
6+
7+
location, err := tpgresource.GetRegion(d, config)
8+
if err != nil {
9+
return nil, err
10+
}
11+
12+
// createNotebookExecutionJobRequest does not exist in update requests
13+
if obj["createNotebookExecutionJobRequest"] == nil {
14+
return obj, nil
15+
}
16+
17+
jobRequest, ok := obj["createNotebookExecutionJobRequest"].(map[string]interface{})
18+
if !ok {
19+
return nil, fmt.Errorf("createNotebookExecutionJobRequest is not of type map[string]interface{} or is nil")
20+
}
21+
jobRequest["parent"] = fmt.Sprintf("projects/%s/locations/%s", project, location)
22+
23+
return obj, nil

‎mmv1/templates/terraform/examples/colab_schedule_basic.tf.tmpl

+87
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,87 @@
1+
resource "google_colab_runtime_template" "my_runtime_template" {
2+
name = "{{index $.Vars "runtime_template_name"}}"
3+
display_name = "Runtime template"
4+
location = "us-central1"
5+
6+
machine_spec {
7+
machine_type = "e2-standard-4"
8+
}
9+
10+
network_spec {
11+
enable_internet_access = true
12+
}
13+
}
14+
15+
resource "google_storage_bucket" "output_bucket" {
16+
name = "{{index $.Vars "bucket"}}"
17+
location = "US"
18+
force_destroy = true
19+
uniform_bucket_level_access = true
20+
}
21+
22+
resource "google_storage_bucket_object" "notebook" {
23+
name = "hello_world.ipynb"
24+
bucket = google_storage_bucket.output_bucket.name
25+
content = <<EOF
26+
{
27+
"cells": [
28+
{
29+
"cell_type": "code",
30+
"execution_count": null,
31+
"metadata": {},
32+
"outputs": [],
33+
"source": [
34+
"print(\"Hello, World!\")"
35+
]
36+
}
37+
],
38+
"metadata": {
39+
"kernelspec": {
40+
"display_name": "Python 3",
41+
"language": "python",
42+
"name": "python3"
43+
},
44+
"language_info": {
45+
"codemirror_mode": {
46+
"name": "ipython",
47+
"version": 3
48+
},
49+
"file_extension": ".py",
50+
"mimetype": "text/x-python",
51+
"name": "python",
52+
"nbconvert_exporter": "python",
53+
"pygments_lexer": "ipython3",
54+
"version": "3.8.5"
55+
}
56+
},
57+
"nbformat": 4,
58+
"nbformat_minor": 4
59+
}
60+
EOF
61+
}
62+
63+
resource "google_colab_schedule" "{{$.PrimaryResourceId}}" {
64+
display_name = "{{index $.Vars "display_name"}}"
65+
location = "{{index $.TestEnvVars "location"}}"
66+
max_concurrent_run_count = 2
67+
cron = "TZ=America/Los_Angeles * * * * *"
68+
69+
create_notebook_execution_job_request {
70+
notebook_execution_job {
71+
display_name = "Notebook execution"
72+
gcs_notebook_source {
73+
uri = "gs://${google_storage_bucket_object.notebook.bucket}/${google_storage_bucket_object.notebook.name}"
74+
generation = google_storage_bucket_object.notebook.generation
75+
}
76+
77+
notebook_runtime_template_resource_name = "projects/${google_colab_runtime_template.my_runtime_template.project}/locations/${google_colab_runtime_template.my_runtime_template.location}/notebookRuntimeTemplates/${google_colab_runtime_template.my_runtime_template.name}"
78+
gcs_output_uri = "gs://${google_storage_bucket.output_bucket.name}"
79+
service_account = "{{index $.TestEnvVars "service_account"}}"
80+
}
81+
}
82+
83+
depends_on = [
84+
google_colab_runtime_template.my_runtime_template,
85+
google_storage_bucket.output_bucket,
86+
]
87+
}

0 commit comments

Comments
 (0)