Implement remote state managment changes (#380)

* Implement remote stage managment

Signed-off-by: Piotr Pawłowski <[email protected]>

---------

Signed-off-by: Piotr Pawłowski <[email protected]>

Author: pawloch00

pyproject.toml

@@ -40,6 +40,7 @@ dependencies = [
   "google-api-core",
   "packaging",
   "google-cloud-filestore",
+  "google-cloud-storage"
 ]
 
 [project.urls]
@@ -65,7 +66,7 @@ dev = [
 ]
 
 [tool.setuptools]
-packages = ["xpk", "xpk.parser", "xpk.core", "xpk.commands", "xpk.api", "xpk.templates", "xpk.utils", "xpk.core.blueprint", "xpk.core.workload_decorators"]
+packages = ["xpk", "xpk.parser", "xpk.core", "xpk.commands", "xpk.api", "xpk.templates", "xpk.utils", "xpk.core.blueprint", "xpk.core.remote_state", "xpk.core.workload_decorators"]
 package-dir = {"" = "src"}
 package-data = {"xpk.api" = ["storage_crd.yaml"], "xpk.templates" = ["storage.yaml"]}
 

src/xpk/commands/cluster_gcluster.py

@@ -16,13 +16,16 @@
 
 import os
 
+from ..core.remote_state.remote_state_client import RemoteStateClient
+from ..core.remote_state.fuse_remote_state import FuseStateClient
 from ..core.blueprint.blueprint_generator import (
     BlueprintGenerator,
     BlueprintGeneratorOutput,
     a3mega_device_type,
     a3ultra_device_type,
     supported_device_types,
 )
+from ..core.commands import run_command_for_value
 from ..core.capacity import get_capacity_type
 from ..core.docker_manager import DockerManager
 from ..core.gcloud_context import zone_to_region
@@ -49,13 +52,22 @@ def cluster_create(args) -> None:
   """
   check_gcloud_authenticated()
   prepare_directories()
-  gcm = prepare_gcluster_manager()
   region = zone_to_region(args.zone)
 
   # unique_name uses shortened hash string, so still name collision is possible
   unique_name = get_unique_name(args.project, region, args.cluster)
   # prefix is to prevent name collisions for blueprints and also deployments by storing them in prefix directory. Ex.: blueprints/{prefix}/cluster_name_hash
   prefix = get_prefix_path(args.project, region)
+  remote_state_client = None
+  if args.cluster_state_gcs_bucket is not None:
+    remote_state_client = FuseStateClient(
+        bucket=args.cluster_state_gcs_bucket,
+        state_directory=os.path.join(blueprints_path, prefix, unique_name),
+        prefix=prefix,
+        cluster=args.cluster,
+        deployment_name=unique_name,
+    )
+  gcm = prepare_gcluster_manager(remote_state_client)
 
   bp = generate_blueprint(blueprint_name=unique_name, args=args, prefix=prefix)
 
@@ -70,6 +82,8 @@ def cluster_create(args) -> None:
       deployment_name=unique_name,
       prefix=prefix,
   )
+  if args.cluster_state_gcs_bucket is not None:
+    gcm.upload_state()
 
   set_cluster_command_code = set_cluster_command(args)
   if set_cluster_command_code != 0:
@@ -89,15 +103,42 @@ def cluster_delete(args) -> None:
   """
   check_gcloud_authenticated()
   prepare_directories()
-  gcm = prepare_gcluster_manager()
   region = zone_to_region(args.zone)
+  unique_name = get_unique_name(args.project, region, args.cluster)
+  # prefix is to prevent name collisions for blueprints and also deployments by storing them in prefix directory. Ex.: blueprints/{prefix}/cluster_name_hash
+  prefix = get_prefix_path(args.project, region)
+  remote_state_client = None
+  if args.cluster_state_gcs_bucket is not None:
+    remote_state_client = FuseStateClient(
+        bucket=args.cluster_state_gcs_bucket,
+        state_directory=os.path.join(blueprints_path, prefix, unique_name),
+        prefix=prefix,
+        cluster=args.cluster,
+        deployment_name=unique_name,
+    )
+  gcm = prepare_gcluster_manager(remote_state_client)
 
   # unique_name uses shortened hash string, so still name collision is possible
   unique_name = get_unique_name(args.project, region, args.cluster)
   # prefix is to prevent name collisions for blueprints and also deployments by storing them in prefix directory. Ex.: blueprints/{prefix}/cluster_name_hash
-  prefix_path = get_prefix_path(args.project, region)
+  prefix = get_prefix_path(args.project, region)
+  if args.cluster_state_gcs_bucket is not None:
+    gcm.download_state()
+
+    bp = BlueprintGeneratorOutput(
+        blueprint_file=os.path.join(blueprints_path, prefix, unique_name)
+        + '.yaml',
+        blueprint_dependencies=os.path.join(
+            blueprints_path, prefix, unique_name
+        ),
+    )
 
-  gcm.destroy_deployment(deployment_name=unique_name, prefix=prefix_path)
+    gcm.stage_files(
+        blueprint_file=bp.blueprint_file,
+        blueprint_dependencies=bp.blueprint_dependencies,
+        prefix=prefix,
+    )
+  gcm.destroy_deployment(deployment_name=unique_name, prefix=prefix)
 
   xpk_exit(0)
 
@@ -140,18 +181,35 @@ def check_gcloud_authenticated():
     xpk_exit(1)
 
 
-def prepare_gcluster_manager() -> GclusterManager:
+def prepare_gcluster_manager(
+    remote_state_client: RemoteStateClient | None,
+) -> GclusterManager:
   dm = DockerManager(
       working_dir=gcluster_working_dir, gcloud_cfg_path=gcloud_cfg_path
   )
   dm.initialize()
-  return GclusterManager(gcluster_command_runner=dm)
+  return GclusterManager(
+      gcluster_command_runner=dm, remote_state_client=remote_state_client
+  )
 
 
 def prepare_blueprint_generator() -> BlueprintGenerator:
   return BlueprintGenerator(storage_path=blueprints_path)
 
 
+def validate_state_gcs_bucket(args):
+  bucket_validate_cmd = (
+      f'gcloud storage buckets describe gs://{args.cluster_state_gcs_bucket}'
+  )
+  err_code, _ = run_command_for_value(
+      bucket_validate_cmd,
+      'Validate remote state bucket existence.',
+      global_args=args,
+  )
+  if err_code != 0:
+    xpk_exit(err_code)
+
+
 def generate_blueprint(
     blueprint_name, args, prefix=None
 ) -> BlueprintGeneratorOutput:
@@ -162,6 +220,9 @@ def generate_blueprint(
 
   bpg = prepare_blueprint_generator()
 
+  if args.cluster_state_gcs_bucket is not None:
+    validate_state_gcs_bucket(args)
+
   if args.device_type in supported_device_types:
     if args.device_type == a3mega_device_type:
       num_nodes = args.num_nodes if not args.num_nodes is None else 2
@@ -178,6 +239,7 @@ def generate_blueprint(
           capacity_type=capacity_type,
           system_node_pool_machine_type=args.default_pool_cpu_machine_type,
           system_node_pool_min_node_count=args.default_pool_cpu_num_nodes,
+          gcs_bucket=args.cluster_state_gcs_bucket,
       )
     if args.device_type == a3ultra_device_type:
       num_nodes = args.num_nodes if not args.num_nodes is None else 2
@@ -195,5 +257,6 @@ def generate_blueprint(
           capacity_type=capacity_type,
           system_node_pool_machine_type=args.default_pool_cpu_machine_type,
           system_node_pool_min_node_count=args.default_pool_cpu_num_nodes,
+          gcs_bucket=args.cluster_state_gcs_bucket,
       )
   return None

src/xpk/core/blueprint/blueprint_generator.py

@@ -641,7 +641,7 @@ def _getblock_terraform_backend(
     }
 
   def _get_terraforrm_backend_full_prefix(self, prefix: str = "") -> str:
-    return f"xpk_terraform_state/{prefix}"
+    return f"xpk_terraform_state/{prefix}/tfstate/"
 
   def _save_blueprint_to_file(
       self, blueprint_name: str, xpk_blueprint: Blueprint, prefix: str = ""

src/xpk/core/gcluster_manager.py

@@ -15,8 +15,8 @@
 """
 
 from .docker_manager import CommandRunner
-from ..utils.console import xpk_print
-
+from ..utils.console import xpk_exit, xpk_print
+from .remote_state.remote_state_client import RemoteStateClient
 
 xpk_gcloud_cfg_path = '~/gcloud/cfg'
 xpk_deployment_dir = '/deployment'
@@ -44,8 +44,10 @@ class GclusterManager:
   def __init__(
       self,
       gcluster_command_runner: CommandRunner,
+      remote_state_client: RemoteStateClient | None,
   ) -> None:
     self.gcluster_command_runner = gcluster_command_runner
+    self.remote_state_client = remote_state_client
 
   def _run_create_deployment_cmd(
       self, blueprint_container_path: str, prefix: str = ''
@@ -156,3 +158,19 @@ def stage_files(
     xpk_print('Staging blueprint completed!')
     xpk_print(f"File path in gcluster's working directory: {staged_blueprint}")
     return staged_blueprint
+
+  def upload_state(self) -> None:
+    xpk_print('Uploading state.')
+    if self.remote_state_client is None:
+      xpk_print('No remote state defined')
+      xpk_exit(1)
+    self.remote_state_client.upload_state()
+
+  def download_state(self) -> None:
+    if self.remote_state_client is None:
+      xpk_print('No remote state defined')
+      xpk_exit(1)
+
+    if self.remote_state_client.check_remote_state_exists():
+      self.remote_state_client.download_state()
+    xpk_print('Remote state not found.')

src/xpk/core/remote_state/__init__.py

@@ -0,0 +1,15 @@
+"""
+Copyright 2025 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""

src/xpk/core/remote_state/fuse_remote_state.py

@@ -0,0 +1,99 @@
+"""
+Copyright 2025 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+from .remote_state_client import RemoteStateClient
+from ...utils.gcs_utils import upload_directory_to_gcs, check_file_exists, download_bucket_to_dir, upload_file_to_gcs
+from ...utils.console import xpk_print
+from google.cloud.storage import Client
+import os
+
+
+class FuseStateClient(RemoteStateClient):
+  """FuseStateClient is a class for managing remote xpk state stored in GCS Fuse."""
+
+  def __init__(
+      self,
+      bucket: str,
+      state_directory: str,
+      cluster: str,
+      deployment_name: str,
+      prefix: str,
+  ) -> None:
+    self.bucket = bucket
+    self.state_dir = state_directory
+    self.storage_client = Client()
+    self.cluster = cluster
+    self.prefix = prefix
+    self.deployment_name = deployment_name
+
+  def _get_bucket_path(self) -> str:
+    return (
+        f'xpk_terraform_state/{self.prefix}/blueprints/{self.deployment_name}/'
+    )
+
+  def _get_bucket_path_blueprint(self) -> str:
+    return f'xpk_terraform_state/{self.prefix}/blueprints/'
+
+  def _get_deployment_filename(self) -> str:
+    return f'{self.deployment_name}.yaml'
+
+  def _get_blueprint_path(self) -> str:
+    blueprint_dir = '/'.join(self.state_dir.split('/')[:-1])
+    return os.path.join(blueprint_dir, self.deployment_name) + '.yaml'
+
+  def upload_state(self) -> None:
+    xpk_print(
+        f'Uploading dependecies from directory {self.state_dir} to bucket:'
+        f' {self.bucket}. Path within bucket is: {self._get_bucket_path()}'
+    )
+    upload_directory_to_gcs(
+        storage_client=self.storage_client,
+        bucket_name=self.bucket,
+        bucket_path=self._get_bucket_path(),
+        source_directory=self.state_dir,
+    )
+    blueprint_bucket_path = (
+        self._get_bucket_path_blueprint() + self._get_deployment_filename()
+    )
+    xpk_print(
+        f'Uploading blueprint file: {self._get_blueprint_path()} to bucket'
+        f' {self.bucket}. Path within bucket is: {blueprint_bucket_path}'
+    )
+    upload_file_to_gcs(
+        storage_client=self.storage_client,
+        bucket_name=self.bucket,
+        bucket_path=blueprint_bucket_path,
+        file=self._get_blueprint_path(),
+    )
+
+  def download_state(self) -> None:
+    xpk_print(
+        f'Downloading from bucket: {self.bucket}, from path:'
+        f' {self._get_bucket_path()} to directory: {self.state_dir}'
+    )
+    download_bucket_to_dir(
+        self.storage_client,
+        self.bucket,
+        self._get_bucket_path(),
+        destination_directory=self.state_dir,
+    )
+
+  def check_remote_state_exists(self) -> bool:
+    return check_file_exists(
+        self.storage_client,
+        self.bucket,
+        self._get_bucket_path_blueprint() + self._get_deployment_filename(),
+    )

src/xpk/core/remote_state/remote_state_client.py

@@ -0,0 +1,38 @@
+"""
+Copyright 2025 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+from abc import ABC, abstractmethod
+
+
+class RemoteStateClient(ABC):
+  """This is a base class that defines methods a class for managing remote cluster state.
+  Args:
+      ABC (_type_): _description_
+  """
+
+  @abstractmethod
+  def upload_state(self) -> None:
+    """Upload state to remote storage"""
+    return None
+
+  @abstractmethod
+  def download_state(self) -> None:
+    """Download state from remote storage"""
+    return None
+
+  @abstractmethod
+  def check_remote_state_exists(self) -> bool:
+    return False

src/xpk/core/tests/data/a3_ultra.yaml

@@ -22,7 +22,7 @@ terraform_backend_defaults:
   type: gcs
   configuration:
     bucket: test-bucket
-    prefix: xpk_terraform_state/testdir
+    prefix: xpk_terraform_state/testdir/tfstate/
 
 deployment_groups:
 - !DeploymentGroup

src/xpk/core/tests/integration/test_gcluster.py

@@ -82,7 +82,7 @@ def test_create_deployment():
   assert os.path.exists(blueprint_test_path)
 
   gcluster_manager = GclusterManager(
-      gcluster_command_runner=docker_manager,
+      gcluster_command_runner=docker_manager, remote_state_client=None
   )
 
   staged_bp_path = gcluster_manager.stage_files(