Merge pull request GoogleCloudPlatform#4539 from alyssa-sm/reduce_zipball

Separate controller and compute zip files

Author: alyssa-sm

community/modules/internal/slurm-gcp/instance_template/files/startup_sh_unlinted

@@ -76,7 +76,15 @@ fi
 
 SCRIPTS_ZIP="$HOME/slurm-gcp-scripts.zip"
 export CLOUDSDK_CORE_UNIVERSE_DOMAIN="$UNIVERSE_DOMAIN"
-until gcloud storage cp "$BUCKET/slurm-gcp-devel.zip" "$SCRIPTS_ZIP"; do
+
+INSTANCE_ROLE="$($CURL $URL/instance/attributes/slurm_instance_role)"
+
+if [ "$INSTANCE_ROLE" == "controller" ]; then
+    DEVEL_ZIP="slurm-gcp-devel-controller.zip"
+else
+    DEVEL_ZIP="slurm-gcp-devel.zip"
+fi
+until gcloud storage cp "$BUCKET/$DEVEL_ZIP" "$SCRIPTS_ZIP"; do
 	echo "WARN: Could not download SlurmGCP scripts, retrying in 5 seconds."
         # Remove marker used to determine if gcloud is being used in a GCE VM.
         # This can get mistakenly set to False in some cases.

community/modules/scheduler/schedmd-slurm-gcp-v6-controller/modules/slurm_files/README.md

@@ -43,6 +43,7 @@ No modules.
 | [google_storage_bucket_object.config](https://registry.terraform.io/providers/hashicorp/google/latest/docs/resources/storage_bucket_object) | resource |
 | [google_storage_bucket_object.controller_startup_scripts](https://registry.terraform.io/providers/hashicorp/google/latest/docs/resources/storage_bucket_object) | resource |
 | [google_storage_bucket_object.devel](https://registry.terraform.io/providers/hashicorp/google/latest/docs/resources/storage_bucket_object) | resource |
+| [google_storage_bucket_object.devel_compute](https://registry.terraform.io/providers/hashicorp/google/latest/docs/resources/storage_bucket_object) | resource |
 | [google_storage_bucket_object.epilog_scripts](https://registry.terraform.io/providers/hashicorp/google/latest/docs/resources/storage_bucket_object) | resource |
 | [google_storage_bucket_object.nodeset_config](https://registry.terraform.io/providers/hashicorp/google/latest/docs/resources/storage_bucket_object) | resource |
 | [google_storage_bucket_object.nodeset_dyn_config](https://registry.terraform.io/providers/hashicorp/google/latest/docs/resources/storage_bucket_object) | resource |
@@ -52,7 +53,8 @@ No modules.
 | [google_storage_bucket_object.task_epilog_scripts](https://registry.terraform.io/providers/hashicorp/google/latest/docs/resources/storage_bucket_object) | resource |
 | [google_storage_bucket_object.task_prolog_scripts](https://registry.terraform.io/providers/hashicorp/google/latest/docs/resources/storage_bucket_object) | resource |
 | [random_uuid.cluster_id](https://registry.terraform.io/providers/hashicorp/random/latest/docs/resources/uuid) | resource |
-| [archive_file.slurm_gcp_devel_zip](https://registry.terraform.io/providers/hashicorp/archive/latest/docs/data-sources/file) | data source |
+| [archive_file.slurm_gcp_devel_compute_zip](https://registry.terraform.io/providers/hashicorp/archive/latest/docs/data-sources/file) | data source |
+| [archive_file.slurm_gcp_devel_controller_zip](https://registry.terraform.io/providers/hashicorp/archive/latest/docs/data-sources/file) | data source |
 | [google_storage_bucket.this](https://registry.terraform.io/providers/hashicorp/google/latest/docs/data-sources/storage_bucket) | data source |
 | [local_file.chs_gpu_health_check](https://registry.terraform.io/providers/hashicorp/local/latest/docs/data-sources/file) | data source |
 | [local_file.external_epilog](https://registry.terraform.io/providers/hashicorp/local/latest/docs/data-sources/file) | data source |

community/modules/scheduler/schedmd-slurm-gcp-v6-controller/modules/slurm_files/main.tf

@@ -169,29 +169,92 @@ resource "google_storage_bucket_object" "nodeset_tpu_config" {
 locals {
   build_dir = abspath("${path.module}/build")
 
-  slurm_gcp_devel_zip        = "slurm-gcp-devel.zip"
-  slurm_gcp_devel_zip_bucket = format("%s/%s", local.bucket_dir, local.slurm_gcp_devel_zip)
+  slurm_gcp_devel_controller_zip     = "slurm-gcp-devel-controller.zip"
+  slurm_gcp_devel_compute_zip        = "slurm-gcp-devel.zip"
+  slurm_gcp_devel_zip_bucket         = format("%s/%s", local.bucket_dir, local.slurm_gcp_devel_controller_zip)
+  slurm_gcp_devel_compute_zip_bucket = format("%s/%s", local.bucket_dir, local.slurm_gcp_devel_compute_zip)
+
+  controller_files = [
+    "tools/gpu-test",
+    "tools/task-epilog",
+    "tools/task-prolog",
+    "conf.py",
+    "file_cache.py",
+    "get_tpu_vmcount.py",
+    "job_submit.lua.tpl",
+    "load_bq.py",
+    "local_pubsub.py",
+    "mig_flex.py",
+    "resume_wrapper.sh",
+    "resume.py",
+    "setup_network_storage.py",
+    "setup.py",
+    "slurmsync.py",
+    "sort_nodes.py",
+    "suspend_wrapper.sh",
+    "suspend.py",
+    "tpu.py",
+    "util.py",
+    "watch_delete_vm_op.py",
+  ]
+
+  compute_files = [
+    "tools/gpu-test",
+    "tools/task-epilog",
+    "tools/task-prolog",
+    "file_cache.py",
+    "get_tpu_vmcount.py",
+    "job_submit.lua.tpl",
+    "local_pubsub.py",
+    "mig_flex.py",
+    "setup_network_storage.py",
+    "setup.py",
+    "slurmsync.py",
+    "sort_nodes.py",
+    "suspend.py",
+    "tpu.py",
+    "util.py",
+    "watch_delete_vm_op.py",
+  ]
 }
 
-data "archive_file" "slurm_gcp_devel_zip" {
-  output_path = "${local.build_dir}/${local.slurm_gcp_devel_zip}"
+data "archive_file" "slurm_gcp_devel_controller_zip" {
+  output_path = "${local.build_dir}/${local.slurm_gcp_devel_controller_zip}"
   type        = "zip"
-  source_dir  = local.scripts_dir
 
-  excludes = flatten([
-    fileset(local.scripts_dir, "tests/**"),
-    # TODO: consider removing (including nested) __pycache__ and all .* files
-    # Though it only affects developers
-  ])
+  dynamic "source" {
+    for_each = local.controller_files
+    content {
+      content  = file("${local.scripts_dir}/${source.value}")
+      filename = source.value
+    }
+  }
+}
+
+data "archive_file" "slurm_gcp_devel_compute_zip" {
+  output_path = "${local.build_dir}/${local.slurm_gcp_devel_compute_zip}"
+  type        = "zip"
 
+  dynamic "source" {
+    for_each = local.compute_files
+    content {
+      content  = file("${local.scripts_dir}/${source.value}")
+      filename = source.value
+    }
+  }
 }
 
 resource "google_storage_bucket_object" "devel" {
   bucket = var.bucket_name
   name   = local.slurm_gcp_devel_zip_bucket
-  source = data.archive_file.slurm_gcp_devel_zip.output_path
+  source = data.archive_file.slurm_gcp_devel_controller_zip.output_path
 }
 
+resource "google_storage_bucket_object" "devel_compute" {
+  bucket = var.bucket_name
+  name   = local.slurm_gcp_devel_compute_zip_bucket
+  source = data.archive_file.slurm_gcp_devel_compute_zip.output_path
+}
 
 ###########
 # SCRIPTS #