Skip to content

Commit

Permalink
Auto-upgrade for colab runtimes (#12931) (#21214)
Browse files Browse the repository at this point in the history
[upstream:b79356dcaaa3925403eade2599558778d385097f]

Signed-off-by: Modular Magician <magic-modules@google.com>
  • Loading branch information
modular-magician authored Feb 3, 2025
1 parent d665f32 commit d1fceac
Show file tree
Hide file tree
Showing 6 changed files with 106 additions and 5 deletions.
3 changes: 3 additions & 0 deletions .changelog/12931.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
```release-note:enhancement
colab: Auto-upgrade for colab runtimes
```
62 changes: 62 additions & 0 deletions google/services/colab/resource_colab_runtime.go
Original file line number Diff line number Diff line change
Expand Up @@ -140,6 +140,21 @@ func ResourceColabRuntime() *schema.Resource {
},
},
},
"expiration_time": {
Type: schema.TypeString,
Computed: true,
Description: `Output only. Timestamp when this NotebookRuntime will be expired.`,
},
"is_upgradable": {
Type: schema.TypeBool,
Computed: true,
Description: `Output only. Checks if the NotebookRuntime is upgradable.`,
},
"notebook_runtime_type": {
Type: schema.TypeString,
Computed: true,
Description: `Output only. The type of the notebook runtime.`,
},
"state": {
Type: schema.TypeString,
Computed: true,
Expand All @@ -151,6 +166,11 @@ func ResourceColabRuntime() *schema.Resource {
Description: `Desired state of the Colab Runtime. Set this field to 'RUNNING' to start the runtime, and 'STOPPED' to stop it.`,
Default: "RUNNING",
},
"auto_upgrade": {
Type: schema.TypeBool,
Optional: true,
Description: `Triggers an upgrade anytime the runtime is started if it is upgradable.`,
},
"project": {
Type: schema.TypeString,
Optional: true,
Expand Down Expand Up @@ -325,6 +345,15 @@ func resourceColabRuntimeRead(d *schema.ResourceData, meta interface{}) error {
if err := d.Set("state", flattenColabRuntimeState(res["state"], d, config)); err != nil {
return fmt.Errorf("Error reading Runtime: %s", err)
}
if err := d.Set("is_upgradable", flattenColabRuntimeIsUpgradable(res["isUpgradable"], d, config)); err != nil {
return fmt.Errorf("Error reading Runtime: %s", err)
}
if err := d.Set("expiration_time", flattenColabRuntimeExpirationTime(res["expirationTime"], d, config)); err != nil {
return fmt.Errorf("Error reading Runtime: %s", err)
}
if err := d.Set("notebook_runtime_type", flattenColabRuntimeNotebookRuntimeType(res["notebookRuntimeType"], d, config)); err != nil {
return fmt.Errorf("Error reading Runtime: %s", err)
}

return nil
}
Expand Down Expand Up @@ -370,6 +399,27 @@ func resourceColabRuntimeUpdate(d *schema.ResourceData, meta interface{}) error
log.Printf("[DEBUG] Colab runtime %q has state %q.", name, state)
}

var upgrade_runtime bool
if d.Get("auto_upgrade").(bool) && d.Get("is_upgradable").(bool) {
upgrade_runtime = true
}

expiration_time_string := d.Get("expiration_time").(string)
expiration_time, err := time.Parse(time.RFC3339Nano, expiration_time_string)
if err != nil {
return err
}

if expiration_time.Before(time.Now()) && d.Get("notebook_runtime_type").(string) == "USER_DEFINED" {
upgrade_runtime = true
}

if upgrade_runtime {
if err := ModifyColabRuntime(config, d, project, billingProject, userAgent, "upgrade"); err != nil {
return err
}
}

return nil
}

Expand Down Expand Up @@ -487,6 +537,18 @@ func flattenColabRuntimeState(v interface{}, d *schema.ResourceData, config *tra
return v
}

func flattenColabRuntimeIsUpgradable(v interface{}, d *schema.ResourceData, config *transport_tpg.Config) interface{} {
return v
}

func flattenColabRuntimeExpirationTime(v interface{}, d *schema.ResourceData, config *transport_tpg.Config) interface{} {
return v
}

func flattenColabRuntimeNotebookRuntimeType(v interface{}, d *schema.ResourceData, config *transport_tpg.Config) interface{} {
return v
}

func expandColabRuntimeNotebookRuntimeTemplateRef(v interface{}, d tpgresource.TerraformResourceData, config *transport_tpg.Config) (interface{}, error) {
l := v.([]interface{})
if len(l) == 0 || l[0] == nil {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -5,14 +5,19 @@ api_service_name: 'aiplatform.googleapis.com'
api_version: 'v1'
api_resource_type_kind: 'NotebookRuntime'
fields:
- field: 'auto_upgrade'
provider_only: true
- field: 'description'
- field: 'desired_state'
provider_only: true
- field: 'display_name'
- field: 'expiration_time'
- field: 'is_upgradable'
- field: 'location'
provider_only: true
- field: 'name'
provider_only: true
- field: 'notebook_runtime_template_ref.notebook_runtime_template'
- field: 'notebook_runtime_type'
- field: 'runtime_user'
- field: 'state'
Original file line number Diff line number Diff line change
Expand Up @@ -170,7 +170,7 @@ func TestAccColabRuntime_colabRuntimeFullExample(t *testing.T) {
ResourceName: "google_colab_runtime.runtime",
ImportState: true,
ImportStateVerify: true,
ImportStateVerifyIgnore: []string{"desired_state", "location", "name"},
ImportStateVerifyIgnore: []string{"auto_upgrade", "desired_state", "location", "name"},
},
},
})
Expand Down Expand Up @@ -235,6 +235,8 @@ resource "google_colab_runtime" "runtime" {
desired_state = "ACTIVE"
auto_upgrade = true
depends_on = [
google_colab_runtime_template.my_template
]
Expand Down
24 changes: 20 additions & 4 deletions google/services/colab/resource_colab_runtime_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ func TestAccColabRuntime_update(t *testing.T) {
ResourceName: "google_colab_runtime.runtime",
ImportState: true,
ImportStateVerify: true,
ImportStateVerifyIgnore: []string{"desired_state", "location", "name"},
ImportStateVerifyIgnore: []string{"desired_state", "location", "name", "auto_upgrade"},
},
{
Config: testAccColabRuntime_no_state(context),
Expand All @@ -40,7 +40,7 @@ func TestAccColabRuntime_update(t *testing.T) {
ResourceName: "google_colab_runtime.runtime",
ImportState: true,
ImportStateVerify: true,
ImportStateVerifyIgnore: []string{"desired_state", "location", "name"},
ImportStateVerifyIgnore: []string{"desired_state", "location", "name", "auto_upgrade"},
},
{
Config: testAccColabRuntime_stopped(context),
Expand All @@ -54,7 +54,21 @@ func TestAccColabRuntime_update(t *testing.T) {
ResourceName: "google_colab_runtime.runtime",
ImportState: true,
ImportStateVerify: true,
ImportStateVerifyIgnore: []string{"desired_state", "location", "name"},
ImportStateVerifyIgnore: []string{"desired_state", "location", "name", "auto_upgrade"},
},
{
Config: testAccColabRuntime_full(context),
ConfigPlanChecks: resource.ConfigPlanChecks{
PreApply: []plancheck.PlanCheck{
plancheck.ExpectResourceAction("google_colab_runtime.runtime", plancheck.ResourceActionUpdate),
},
},
},
{
ResourceName: "google_colab_runtime.runtime",
ImportState: true,
ImportStateVerify: true,
ImportStateVerifyIgnore: []string{"desired_state", "location", "name", "auto_upgrade"},
},
},
})
Expand Down Expand Up @@ -117,7 +131,9 @@ resource "google_colab_runtime" "runtime" {
runtime_user = "gterraformtestuser@gmail.com"
description = "Full runtime"
desired_state = "ACTIVE"
desired_state = "RUNNING"
auto_upgrade = true
depends_on = [
google_colab_runtime_template.my_template
Expand Down
13 changes: 13 additions & 0 deletions website/docs/r/colab_runtime.html.markdown
Original file line number Diff line number Diff line change
Expand Up @@ -174,6 +174,8 @@ resource "google_colab_runtime" "runtime" {
desired_state = "ACTIVE"
auto_upgrade = true
depends_on = [
google_colab_runtime_template.my_template
]
Expand Down Expand Up @@ -219,6 +221,8 @@ The following arguments are supported:

* `desired_state` - (Optional) Desired state of the Colab Runtime. Set this field to `RUNNING` to start the runtime, and `STOPPED` to stop it.

* `auto_upgrade` - (Optional) Triggers an upgrade anytime the runtime is started if it is upgradable.


<a name="nested_notebook_runtime_template_ref"></a>The `notebook_runtime_template_ref` block supports:

Expand All @@ -235,6 +239,15 @@ In addition to the arguments listed above, the following computed attributes are
* `state` -
Output only. The state of the runtime.

* `is_upgradable` -
Output only. Checks if the NotebookRuntime is upgradable.

* `expiration_time` -
Output only. Timestamp when this NotebookRuntime will be expired.

* `notebook_runtime_type` -
Output only. The type of the notebook runtime.


## Timeouts

Expand Down

0 comments on commit d1fceac

Please # to comment.