0.5.2 dev

test/integration/e2e_logging_paths.sh

@@ -27,7 +27,7 @@ readonly LOGGING_BASE="$(dirname "${LOGGING}")"
 declare LOGGING_OVERRIDE
 
 readonly JOB_NAME="log-test"
-readonly JOB_USER="${USER}"
+readonly JOB_USER="${USER:-jupyter}"
 
 # Test a basic job with base logging path
 echo "Subtest #1: Basic logging path"

test/integration/e2e_logging_paths_pattern_tasks.sh

@@ -31,7 +31,7 @@ readonly LOGGING_BASE="$(dirname "${LOGGING}")"
 declare LOGGING_OVERRIDE
 
 readonly JOB_NAME=$(logging_paths_tasks_setup::get_job_name)
-readonly JOB_USER="${USER}"
+readonly JOB_USER="${USER:-jupyter}"
 
 # Set up the tasks file
 logging_paths_tasks_setup::write_tasks_file

test/integration/io_setup.sh

@@ -230,7 +230,7 @@ function io_setup::check_dstat() {
   local dstat_output=$(run_dstat --status '*' --jobs "${job_id}" --full)
 
   echo "  Checking user-id"
-  util::dstat_yaml_assert_field_equal "${dstat_output}" "[0].user-id" "${USER}"
+  util::dstat_yaml_assert_field_equal "${dstat_output}" "[0].user-id" "${USER:-jupyter}"
 
   echo "  Checking logging"
   util::dstat_yaml_assert_field_equal "${dstat_output}" "[0].logging" "${LOGGING}"

test/integration/io_tasks_setup.sh

@@ -119,7 +119,7 @@ function io_tasks_setup::check_dstat() {
     echo "  Check task ${task_id}"
 
     echo "    Checking user-id"
-    util::dstat_yaml_assert_field_equal "${dstat_output}" "[0].user-id" "${USER}"
+    util::dstat_yaml_assert_field_equal "${dstat_output}" "[0].user-id" "${USER:-jupyter}"
 
     echo "    Checking logging"
     util::dstat_yaml_assert_field_equal "${dstat_output}" "[0].logging" "${LOGGING}/${job_id}.${task_id}.log"

test/integration/test_setup.sh

@@ -30,6 +30,9 @@
 # * Provide functions run_dsub, run_dstat, run_ddel which will call a function
 #   with DSUB_PROVIDER-specific default parameters set.
 
+# Set default USER if not already set (needed for Jupyterlab/Docker environments)
+export USER="${USER:-jupyter}"
+
 # If the DSUB_PROVIDER is not set, figure it out from the name of the script.
 #   If the script name is <test>.<provider>.sh, pull out the provider.
 #   If the script name is <test>.sh, use "local".
@@ -89,16 +92,23 @@ function run_dsub() {
 }
 
 function dsub_google-batch() {
-  local location="${LOCATION:-}"
+  # Use REGIONS env var if set, otherwise fall back to LOCATION
+  local location="${LOCATION:-${REGIONS:-}}"
+
+  # Use environment variables for VPC-SC configuration if set
+  local network="${GPU_NETWORK:-global/networks/default}"
+  local subnetwork="${GPU_SUBNETWORK:-regions/us-central1/subnetworks/default}"
+  local service_account="${PET_SA_EMAIL:-}"
 
   dsub \
     --provider google-batch \
     --project "${PROJECT_ID}" \
     ${location:+--location "${location}"} \
     --logging "${LOGGING_OVERRIDE:-${LOGGING}}" \
-    --network "global/networks/default" \
-    --subnetwork "regions/us-central1/subnetworks/default" \
+    --network "${network}" \
+    --subnetwork "${subnetwork}" \
     --use-private-address \
+    ${service_account:+--service-account "${service_account}"} \
     "${@}"
 }
 

test/integration/test_setup_e2e.py

@@ -171,16 +171,27 @@ def dsub_google_batch(dsub_args):
     if val:
       opt_args.append(var[1], val)
 
-  # pyformat: disable
-  return dsub_command.call([
+  # Use environment variables for VPC-SC configuration if set
+  network = os.environ.get("GPU_NETWORK", "global/networks/default")
+  subnetwork = os.environ.get("GPU_SUBNETWORK",
+                               "regions/us-central1/subnetworks/default")
+  location = os.environ.get("LOCATION", os.environ.get("REGIONS", "us-central1"))
+  service_account = os.environ.get("PET_SA_EMAIL", "")
+
+  args = [
       "--provider", "google-batch",
       "--project", PROJECT_ID,
       "--logging", LOGGING,
-      "--regions", "us-central1",
-      "--network", "global/networks/default",
-      "--subnetwork", "regions/us-central1/subnetworks/default",
+      "--regions", location,
+      "--network", network,
+      "--subnetwork", subnetwork,
       "--use-private-address"
-      ] + opt_args + dsub_args)
+  ]
+  if service_account:
+    args.extend(["--service-account", service_account])
+
+  # pyformat: disable
+  return dsub_command.call(args + opt_args + dsub_args)
   # pyformat: enable