diff --git a/kinetic/backend/execution.py b/kinetic/backend/execution.py index 680665b5..819194bd 100644 --- a/kinetic/backend/execution.py +++ b/kinetic/backend/execution.py @@ -77,6 +77,7 @@ class JobContext: image_uri: Optional[str] = None payload_sha256: Optional[str] = None context_sha256: Optional[str] = None + signed_urls: Optional[dict[str, str]] = None def __post_init__(self): self.bucket_name = build_bucket_name(self.project, self.cluster_name) @@ -203,6 +204,7 @@ def submit_job(self, ctx: JobContext) -> Any: debug=ctx.debug, payload_sha256=ctx.payload_sha256, context_sha256=ctx.context_sha256, + signed_urls=ctx.signed_urls, ) def wait_for_job(self, job: Any, ctx: JobContext) -> None: @@ -261,6 +263,7 @@ def submit_job(self, ctx: JobContext) -> Any: debug=ctx.debug, payload_sha256=ctx.payload_sha256, context_sha256=ctx.context_sha256, + signed_urls=ctx.signed_urls, ) def wait_for_job(self, job: Any, ctx: JobContext) -> None: @@ -575,6 +578,27 @@ def _upload_artifacts(ctx: JobContext) -> bool: project=ctx.project, requirements_content=requirements_content, ) + + # Generate signed URLs for the job (with fallback to legacy GCS access) + signer_email = f"kn-{ctx.cluster_name}-signer@{ctx.project}.iam.gserviceaccount.com" + try: + ctx.signed_urls = storage.generate_job_signed_urls( + bucket_name=ctx.bucket_name, + job_id=ctx.job_id, + project=ctx.project, + signer_sa_email=signer_email, + has_requirements=has_requirements, + debug=ctx.debug, + ) + logging.info("Successfully generated Signed URLs for secure job isolation.") + except Exception as e: + logging.warning( + "Failed to generate Signed URLs (likely due to missing 'signer' GSA or lack of " + "impersonation permissions). Falling back to legacy GCS access. " + "Error: %s", e + ) + ctx.signed_urls = None + return has_requirements diff --git a/kinetic/backend/gke_client.py b/kinetic/backend/gke_client.py index dfaf42ee..c26a2b0e 100644 --- a/kinetic/backend/gke_client.py +++ b/kinetic/backend/gke_client.py @@ -30,6 +30,7 @@ def submit_k8s_job( debug=False, payload_sha256=None, context_sha256=None, + signed_urls=None, ): """Submit a Kubernetes Job to GKE cluster. @@ -45,6 +46,7 @@ def submit_k8s_job( install (prebuilt image mode). payload_sha256: Optional SHA-256 hash of payload.pkl for verification. context_sha256: Optional SHA-256 hash of context.zip for verification. + signed_urls: Optional dict of GCS Signed URLs for job isolation. Returns: kubernetes.client.V1Job object @@ -66,6 +68,7 @@ def submit_k8s_job( debug=debug, payload_sha256=payload_sha256, context_sha256=context_sha256, + signed_urls=signed_urls, ) # Submit job @@ -322,6 +325,7 @@ def _create_job_spec( debug=False, payload_sha256=None, context_sha256=None, + signed_urls=None, ): """Create Kubernetes Job specification. @@ -336,6 +340,7 @@ def _create_job_spec( install (prebuilt image mode). payload_sha256: Optional SHA-256 hash of payload.pkl context_sha256: Optional SHA-256 hash of context.zip + signed_urls: Optional dict of GCS Signed URLs for job isolation. Returns: V1Job object ready for creation @@ -364,16 +369,30 @@ def _create_job_spec( ) # Container arguments: context, payload, result, [requirements] - container_args = [ - "--context-gcs", - f"gs://{bucket_name}/{job_id}/context.zip", - "--payload-gcs", - f"gs://{bucket_name}/{job_id}/payload.pkl", - "--result-gcs", - f"gs://{bucket_name}/{job_id}/result.pkl", - ] - if requirements_uri: - container_args.extend(["--requirements-gcs", requirements_uri]) + if signed_urls: + container_args = [ + "--context-gcs", signed_urls["context_download"], + "--payload-gcs", signed_urls["payload_download"], + "--result-gcs", signed_urls["result_upload"], + ] + if "requirements_download" in signed_urls: + container_args.extend(["--requirements-gcs", signed_urls["requirements_download"]]) + if "debug_ready_upload" in signed_urls: + container_args.extend(["--debug-ready-url", signed_urls["debug_ready_upload"]]) + if "leader_ready_upload" in signed_urls: + container_args.extend(["--leader-ready-upload-url", signed_urls["leader_ready_upload"]]) + if "leader_ready_download" in signed_urls: + container_args.extend(["--leader-ready-download-url", signed_urls["leader_ready_download"]]) + else: + # Legacy GCS URIs + container_args = [ + "--context-gcs", f"gs://{bucket_name}/{job_id}/context.zip", + "--payload-gcs", f"gs://{bucket_name}/{job_id}/payload.pkl", + "--result-gcs", f"gs://{bucket_name}/{job_id}/result.pkl", + ] + if requirements_uri: + container_args.extend(["--requirements-gcs", requirements_uri]) + if payload_sha256: container_args.extend(["--payload-sha256", payload_sha256]) if context_sha256: diff --git a/kinetic/backend/pathways_client.py b/kinetic/backend/pathways_client.py index 32f91c6d..82cc8b2d 100644 --- a/kinetic/backend/pathways_client.py +++ b/kinetic/backend/pathways_client.py @@ -73,6 +73,7 @@ def submit_pathways_job( debug=False, payload_sha256=None, context_sha256=None, + signed_urls=None, ): """Submit a LeaderWorkerSet to GKE cluster. @@ -88,6 +89,7 @@ def submit_pathways_job( install (prebuilt image mode). payload_sha256: Optional SHA-256 hash of payload.pkl for verification. context_sha256: Optional SHA-256 hash of context.zip for verification. + signed_urls: Optional dict of GCS Signed URLs for job isolation. Returns: dict: The created LeaderWorkerSet object @@ -120,6 +122,7 @@ def submit_pathways_job( debug=debug, payload_sha256=payload_sha256, context_sha256=context_sha256, + signed_urls=signed_urls, ) custom_api = _custom_api() @@ -440,6 +443,7 @@ def _create_lws_spec( debug=False, payload_sha256=None, context_sha256=None, + signed_urls=None, ): """Create a LeaderWorkerSet manifest.""" @@ -466,16 +470,31 @@ def _create_lws_spec( entry["value"] = t["value"] tolerations.append(entry) - container_args = [ - "--context-gcs", - f"gs://{bucket_name}/{job_id}/context.zip", - "--payload-gcs", - f"gs://{bucket_name}/{job_id}/payload.pkl", - "--result-gcs", - f"gs://{bucket_name}/{job_id}/result.pkl", - ] - if requirements_uri: - container_args.extend(["--requirements-gcs", requirements_uri]) + # Container arguments: context, payload, result, [requirements] + if signed_urls: + container_args = [ + "--context-gcs", signed_urls["context_download"], + "--payload-gcs", signed_urls["payload_download"], + "--result-gcs", signed_urls["result_upload"], + ] + if "requirements_download" in signed_urls: + container_args.extend(["--requirements-gcs", signed_urls["requirements_download"]]) + if "debug_ready_upload" in signed_urls: + container_args.extend(["--debug-ready-url", signed_urls["debug_ready_upload"]]) + if "leader_ready_upload" in signed_urls: + container_args.extend(["--leader-ready-upload-url", signed_urls["leader_ready_upload"]]) + if "leader_ready_download" in signed_urls: + container_args.extend(["--leader-ready-download-url", signed_urls["leader_ready_download"]]) + else: + # Legacy GCS URIs + container_args = [ + "--context-gcs", f"gs://{bucket_name}/{job_id}/context.zip", + "--payload-gcs", f"gs://{bucket_name}/{job_id}/payload.pkl", + "--result-gcs", f"gs://{bucket_name}/{job_id}/result.pkl", + ] + if requirements_uri: + container_args.extend(["--requirements-gcs", requirements_uri]) + if payload_sha256: container_args.extend(["--payload-sha256", payload_sha256]) if context_sha256: diff --git a/kinetic/cli/infra/program.py b/kinetic/cli/infra/program.py index fa15151d..b92067d1 100644 --- a/kinetic/cli/infra/program.py +++ b/kinetic/cli/infra/program.py @@ -194,14 +194,23 @@ def _create_service_accounts( repo: gcp.artifactregistry.Repository, jobs_bucket: gcp.storage.Bucket, builds_bucket: gcp.storage.Bucket, + deployer_email: pulumi.Input[str], enabled_apis: list[gcp.projects.Service], -) -> tuple[gcp.serviceaccount.Account, gcp.serviceaccount.Account]: - """Create node and build service accounts with their IAM bindings.""" +) -> tuple[ + gcp.serviceaccount.Account, + gcp.serviceaccount.Account, + gcp.serviceaccount.Account, + gcp.serviceaccount.Account, +]: + """Create service accounts with their IAM bindings. + + Returns: + Tuple of (node_sa, worker_sa, signer_sa, build_sa) + """ api_deps = pulumi.ResourceOptions(depends_on=enabled_apis) bucket_pairs = [("jobs", jobs_bucket), ("builds", builds_bucket)] - # Node SA — used by GKE workload pods (GCS, logging, monitoring, - # pulling images from AR). + # 1. Node SA — used by GKE nodes (NO GCS bucket access, only logging/monitoring/AR). node_sa = gcp.serviceaccount.Account( "kinetic-node-sa", account_id=f"kn-{cluster_name}-nodes", @@ -218,14 +227,83 @@ def _create_service_accounts( "roles/monitoring.metricWriter", "roles/container.defaultNodeServiceAccount", ], - bucket_pairs, + [], # No GCS bucket access for nodes repo, ar_location, "roles/artifactregistry.reader", enabled_apis, ) - # Build SA — used as the Cloud Build execution SA. + # 2. Worker SA — used by GKE workload pods (restricted GCS access to cache). + worker_sa = gcp.serviceaccount.Account( + "kinetic-worker-sa", + account_id=f"kn-{cluster_name}-workers", + display_name=f"kinetic {cluster_name} worker SA", + project=project_id, + opts=api_deps, + ) + # Basic project roles for workers (logging/monitoring) + for role in ["roles/logging.logWriter", "roles/monitoring.metricWriter"]: + gcp.projects.IAMMember( + f"worker-sa-{role.split('/')[-1]}", + project=project_id, + role=role, + member=worker_sa.email.apply(lambda e: f"serviceAccount:{e}"), + opts=api_deps, + ) + # Restricted GCS access: read-only on data-cache and data-markers + gcp.storage.BucketIAMMember( + "worker-sa-storage-jobs-viewer", + bucket=jobs_bucket.name, + role="roles/storage.objectViewer", + member=worker_sa.email.apply(lambda e: f"serviceAccount:{e}"), + condition=gcp.storage.BucketIAMMemberConditionArgs( + title="restrict-to-cache", + expression=pulumi.Output.format( + "resource.name.startsWith('projects/_/buckets/{0}/objects/default/data-cache/') || " + "resource.name.startsWith('projects/_/buckets/{0}/objects/default/data-markers/')", + jobs_bucket.name, + ), + ), + ) + # Legacy bucket reader (required for FUSE mount to get bucket metadata) + gcp.storage.BucketIAMMember( + "worker-sa-storage-jobs-bucket-reader", + bucket=jobs_bucket.name, + role="roles/storage.legacyBucketReader", + member=worker_sa.email.apply(lambda e: f"serviceAccount:{e}"), + ) + + # 3. Signer SA — used by client to sign GCS URLs. + signer_sa = gcp.serviceaccount.Account( + "kinetic-signer-sa", + account_id=f"kn-{cluster_name}-signer", + display_name=f"kinetic {cluster_name} signer SA", + project=project_id, + opts=api_deps, + ) + # Full GCS admin on jobs bucket + gcp.storage.BucketIAMMember( + "signer-sa-storage-jobs-admin", + bucket=jobs_bucket.name, + role="roles/storage.objectAdmin", + member=signer_sa.email.apply(lambda e: f"serviceAccount:{e}"), + ) + gcp.storage.BucketIAMMember( + "signer-sa-storage-jobs-bucket-reader", + bucket=jobs_bucket.name, + role="roles/storage.legacyBucketReader", + member=signer_sa.email.apply(lambda e: f"serviceAccount:{e}"), + ) + # Allow deployer to impersonate signer GSA + gcp.serviceaccount.IAMMember( + "deployer-impersonate-signer", + service_account_id=signer_sa.name, + role="roles/iam.serviceAccountTokenCreator", + member=pulumi.Output.format("user:{0}", deployer_email), + ) + + # 4. Build SA — used as the Cloud Build execution SA (unchanged). build_sa = gcp.serviceaccount.Account( "kinetic-build-sa", account_id=f"kn-{cluster_name}-builds", @@ -245,7 +323,7 @@ def _create_service_accounts( enabled_apis, ) - return node_sa, build_sa + return node_sa, worker_sa, signer_sa, build_sa def _create_firewall_cleanup( @@ -387,7 +465,7 @@ def _create_gke_cluster( def _create_k8s_resources( cluster: gcp.container.Cluster, - node_sa: gcp.serviceaccount.Account, + worker_sa: gcp.serviceaccount.Account, project_id: str, node_pools: list[NodePoolConfig], ) -> None: @@ -403,10 +481,10 @@ def _create_k8s_resources( ) # Workload Identity binding — allow the kinetic KSA to impersonate the - # node GSA. + # worker GSA. gcp.serviceaccount.IAMMember( "wif-kinetic-ksa", - service_account_id=node_sa.name, + service_account_id=worker_sa.name, role="roles/iam.workloadIdentityUser", member=pulumi.Output.format( "serviceAccount:{0}.svc.id.goog[default/{1}]", @@ -423,7 +501,7 @@ def _create_k8s_resources( name=KINETIC_KSA_NAME, namespace="default", annotations={ - "iam.gke.io/gcp-service-account": node_sa.email, + "iam.gke.io/gcp-service-account": worker_sa.email, }, ), opts=pulumi.ResourceOptions(provider=k8s_provider), @@ -491,6 +569,8 @@ def _export_stack_outputs( zone: str, cluster: gcp.container.Cluster, node_sa: gcp.serviceaccount.Account, + worker_sa: gcp.serviceaccount.Account, + signer_sa: gcp.serviceaccount.Account, repo: gcp.artifactregistry.Repository, ar_location: str, cluster_name: str, @@ -503,6 +583,8 @@ def _export_stack_outputs( pulumi.export("cluster_name", cluster.name) pulumi.export("cluster_endpoint", cluster.endpoint) pulumi.export("node_sa_email", node_sa.email) + pulumi.export("worker_sa_email", worker_sa.email) + pulumi.export("signer_sa_email", signer_sa.email) pulumi.export("force_destroy", force_destroy) pulumi.export( "ar_registry", @@ -565,6 +647,10 @@ def pulumi_program() -> None: enabled_apis = _enable_apis(project_id) + # Get deployer email for signer GSA impersonation + client_config = gcp.organizations.get_client_config() + deployer_email = client_config.email + repo = gcp.artifactregistry.Repository( "kinetic-repo", repository_id=f"kn-{cluster_name}", @@ -584,13 +670,14 @@ def pulumi_program() -> None: config.force_destroy, ) - node_sa, _build_sa = _create_service_accounts( + node_sa, worker_sa, signer_sa, _build_sa = _create_service_accounts( project_id, cluster_name, ar_location, repo, jobs_bucket, builds_bucket, + deployer_email, enabled_apis, ) @@ -602,19 +689,20 @@ def pulumi_program() -> None: cluster_name, zone, network, - node_sa, + node_sa, # Nodes still use node_sa enabled_apis, firewall_cleanup, ) - _create_k8s_resources(cluster, node_sa, project_id, config.node_pools) + # K8s resources (WIF KSA) use worker_sa + _create_k8s_resources(cluster, worker_sa, project_id, config.node_pools) pool_entries = _create_accelerator_pools( cluster, config.node_pools, zone, project_id, - node_sa.email, + node_sa.email, # Node pools still use node_sa ) _export_stack_outputs( @@ -622,6 +710,8 @@ def pulumi_program() -> None: zone, cluster, node_sa, + worker_sa, + signer_sa, repo, ar_location, cluster_name, diff --git a/kinetic/cli/infra/program_test.py b/kinetic/cli/infra/program_test.py index 3a0baa24..22e2bcc1 100644 --- a/kinetic/cli/infra/program_test.py +++ b/kinetic/cli/infra/program_test.py @@ -230,5 +230,95 @@ def test_default_node_pool_has_kinetic_label(self): ) +class TestServiceAccountsAndIAM(absltest.TestCase): + """Verify that the split GSA architecture and restricted IAM are created.""" + + def _run_program(self, config=None): + config = config or _make_config() + with ( + mock.patch.object(program, "pulumi") as pulumi_mock, + mock.patch.object(program, "command"), + mock.patch.object(program, "gcp") as gcp_mock, + mock.patch.object(program, "k8s") as k8s_mock, + ): + # Mock client config for deployer email + gcp_mock.organizations.get_client_config.return_value = mock.MagicMock( + email="deployer@google.com" + ) + program.create_program(config)() + return gcp_mock, k8s_mock, pulumi_mock + + def test_creates_four_service_accounts(self): + gcp_mock, _, _ = self._run_program() + + sa_calls = gcp_mock.serviceaccount.Account.call_args_list + # Expect: nodes, workers, signer, builds + self.assertLen(sa_calls, 4) + + sa_names = {call.kwargs["account_id"] for call in sa_calls} + self.assertIn("kn-test-cluster-nodes", sa_names) + self.assertIn("kn-test-cluster-workers", sa_names) + self.assertIn("kn-test-cluster-signer", sa_names) + self.assertIn("kn-test-cluster-builds", sa_names) + + def test_worker_sa_has_restricted_gcs_access(self): + gcp_mock, _, pulumi_mock = self._run_program() + + # Check BucketIAMMember calls + iam_calls = gcp_mock.storage.BucketIAMMember.call_args_list + + # Find calls for the worker SA + worker_calls = [ + c for c in iam_calls + if "worker-sa" in c.args[0] or "workers" in c.args[0] + ] + + # Worker should have read-only (objectViewer) access + viewer_calls = [c for c in worker_calls if c.kwargs.get("role") == "roles/storage.objectViewer"] + self.assertLen(viewer_calls, 1) + + # Verify condition expression via pulumi.Output.format calls + format_calls = pulumi_mock.Output.format.call_args_list + self.assertTrue(any("default/data-cache/" in call.args[0] for call in format_calls)) + + # Worker should NOT have objectAdmin + admin_calls = [c for c in worker_calls if c.kwargs.get("role") == "roles/storage.objectAdmin"] + self.assertEmpty(admin_calls) + + def test_signer_sa_has_admin_access_and_impersonation(self): + gcp_mock, _, pulumi_mock = self._run_program() + + # Signer GSA should have objectAdmin on the jobs bucket (1 call, not 2) + iam_calls = gcp_mock.storage.BucketIAMMember.call_args_list + signer_gcs_calls = [ + c for c in iam_calls + if "signer-sa" in c.args[0] and c.kwargs.get("role") == "roles/storage.objectAdmin" + ] + self.assertLen(signer_gcs_calls, 1) + + # Deployer should have Token Creator on signer GSA + sa_iam_calls = gcp_mock.serviceaccount.IAMMember.call_args_list + impersonation_calls = [ + c for c in sa_iam_calls + if "signer" in c.args[0] and c.kwargs.get("role") == "roles/iam.serviceAccountTokenCreator" + ] + self.assertLen(impersonation_calls, 1) + + # Verify member format call + format_calls = pulumi_mock.Output.format.call_args_list + self.assertTrue(any("user:{0}" in call.args[0] for call in format_calls)) + + def test_node_sa_has_no_gcs_access(self): + gcp_mock, _, _ = self._run_program() + + iam_calls = gcp_mock.storage.BucketIAMMember.call_args_list + node_gcs_calls = [ + c for c in iam_calls + if "node-sa" in c.args[0] + ] + # Nodes SA should have NO GCS bucket bindings now + self.assertEmpty(node_gcs_calls) + + if __name__ == "__main__": absltest.main() diff --git a/kinetic/runner/remote_runner.py b/kinetic/runner/remote_runner.py index abf7dea5..04bb92a2 100644 --- a/kinetic/runner/remote_runner.py +++ b/kinetic/runner/remote_runner.py @@ -17,6 +17,8 @@ import time import traceback import urllib.parse +import urllib.request +import urllib.error import zipfile import cloudpickle @@ -70,6 +72,15 @@ def main(): parser.add_argument( "--context-sha256", help="Expected SHA-256 hash of context" ) + parser.add_argument( + "--debug-ready-url", help="Signed URL to upload debug ready sentinel" + ) + parser.add_argument( + "--leader-ready-upload-url", help="Signed URL to upload leader ready sentinel" + ) + parser.add_argument( + "--leader-ready-download-url", help="Signed URL to download leader ready sentinel" + ) args_parsed, _ = parser.parse_known_args() @@ -173,14 +184,24 @@ def main(): # so there's a single source of truth. Fall back to 5678 (debugpy's # default and VS Code's auto-fill) if the env var is missing. debug_port = int(os.environ.get("KINETIC_DEBUG_PORT", 5678)) - debugger_attached = _start_debug_server(debug_port) + debugger_attached = _start_debug_server(debug_port, args_parsed.debug_ready_url) # Signal workers (if any) that the leader is about to call the # user function, so they can proceed without racing ahead and # hanging on the distributed runtime. - _upload_leader_ready_sentinel() + if args_parsed.leader_ready_upload_url: + _upload_sentinel(args_parsed.leader_ready_upload_url) + else: + _upload_leader_ready_sentinel() elif is_debug_worker: # Pathways worker pod — wait for leader's sentinel before running. - _wait_for_leader_ready_sentinel() + if args_parsed.leader_ready_download_url: + leader_timeout = int( + os.environ.get("KINETIC_DEBUG_WAIT_TIMEOUT", _DEBUG_WAIT_TIMEOUT_DEFAULT) + ) + timeout = leader_timeout + _WORKER_WAIT_BUFFER_SECONDS + _wait_for_sentinel(args_parsed.leader_ready_download_url, timeout) + else: + _wait_for_leader_ready_sentinel() # Execute function and capture result logging.info("Executing %s()", func.__name__) @@ -357,6 +378,43 @@ def _wait_for_leader_ready_sentinel(): ) +def _upload_sentinel(url: str): + """Upload an empty sentinel to a Signed URL.""" + logging.info("Uploading sentinel to: %s", url) + try: + req = urllib.request.Request(url, data=b"", method="PUT") + with urllib.request.urlopen(req) as response: + if response.status not in (200, 201): + logging.warning("Failed to upload sentinel to %s: status %d", url, response.status) + else: + logging.info("Successfully uploaded sentinel to %s", url) + except Exception as e: + logging.warning("Failed to upload sentinel to %s: %s", url, e) + + +def _wait_for_sentinel(url: str, timeout: int): + """Poll a Signed URL until it returns 200 (exists), or time out.""" + poll_interval = 5 + deadline = time.monotonic() + timeout + logging.info("Waiting up to %ds for sentinel at %s", timeout, url) + while time.monotonic() < deadline: + try: + req = urllib.request.Request(url, method="GET") + with urllib.request.urlopen(req) as response: + if response.status == 200: + logging.info("Sentinel is ready, proceeding.") + return + except urllib.error.HTTPError as e: + if e.code == 404: + logging.debug("Sentinel not ready yet (404)...") + else: + logging.warning("Error polling sentinel: %s", e) + except Exception as e: + logging.warning("Error polling sentinel: %s", e) + time.sleep(poll_interval) + raise RuntimeError(f"Sentinel did not appear within {timeout}s at {url}") + + def _install_debugger(): """Install debugpy via uv pip at pod startup.""" logging.info("Installing debugpy...") @@ -379,7 +437,7 @@ def _install_debugger(): _DEBUG_WAIT_TIMEOUT_DEFAULT = 600 -def _start_debug_server(port): +def _start_debug_server(port, debug_ready_url=None): """Start debugpy server and wait for client attachment. Waits up to ``KINETIC_DEBUG_WAIT_TIMEOUT`` seconds (default 600) for @@ -388,6 +446,7 @@ def _start_debug_server(port): Args: port: TCP port for debugpy to listen on. + debug_ready_url: Optional Signed URL to upload debug ready sentinel. Returns: True if a debugger client attached, False if timed out. @@ -399,22 +458,25 @@ def _start_debug_server(port): debugpy.listen(("0.0.0.0", port)) try: - # Signal readiness via a GCS sentinel so the local client can detect it. - # Use env vars set by the pod spec rather than parsing sys.argv. - bucket_name = os.environ.get("GCS_BUCKET") - job_id = os.environ.get("JOB_ID") - if not bucket_name or not job_id: - logging.warning("GCS_BUCKET or JOB_ID not set; skipping debug sentinel.") + if debug_ready_url: + _upload_sentinel(debug_ready_url) else: - blob = storage.Client().bucket(bucket_name).blob(f"{job_id}/.debug_ready") - blob.upload_from_string("") - logging.info( - "Published debugpy GCS sentinel to gs://%s/%s/.debug_ready", - bucket_name, - job_id, - ) - except cloud_exceptions.GoogleCloudError as e: - logging.warning("Failed to publish debug readiness sentinel to GCS: %s", e) + # Signal readiness via a GCS sentinel so the local client can detect it. + # Use env vars set by the pod spec rather than parsing sys.argv. + bucket_name = os.environ.get("GCS_BUCKET") + job_id = os.environ.get("JOB_ID") + if not bucket_name or not job_id: + logging.warning("GCS_BUCKET or JOB_ID not set; skipping debug sentinel.") + else: + blob = storage.Client().bucket(bucket_name).blob(f"{job_id}/.debug_ready") + blob.upload_from_string("") + logging.info( + "Published debugpy GCS sentinel to gs://%s/%s/.debug_ready", + bucket_name, + job_id, + ) + except Exception as e: + logging.warning("Failed to publish debug readiness sentinel: %s", e) logging.info("[DEBUGPY] Ready \u2014 listening on 0.0.0.0:%d", port) @@ -644,13 +706,18 @@ def _download_data( def _download_from_gcs(client, gcs_path, local_path): - """Download file from GCS. + """Download file from GCS or Signed URL. Args: - client: Cloud Storage client - gcs_path: GCS URI (gs://bucket/path) + client: Cloud Storage client (can be None if using Signed URL) + gcs_path: GCS URI (gs://...) or Signed URL (http://... or https://...) local_path: Local file path """ + if gcs_path.startswith("http://") or gcs_path.startswith("https://"): + logging.info("Downloading from Signed URL: %s", gcs_path) + urllib.request.urlretrieve(gcs_path, local_path) + return + # Parse gs://bucket/path format parts = gcs_path.replace("gs://", "").split("/", 1) bucket_name = parts[0] @@ -662,13 +729,23 @@ def _download_from_gcs(client, gcs_path, local_path): def _upload_to_gcs(client, local_path, gcs_path): - """Upload file to GCS. + """Upload file to GCS or Signed URL. Args: - client: Cloud Storage client + client: Cloud Storage client (can be None if using Signed URL) local_path: Local file path - gcs_path: GCS URI (gs://bucket/path) + gcs_path: GCS URI (gs://...) or Signed URL (http://... or https://...) """ + if gcs_path.startswith("http://") or gcs_path.startswith("https://"): + logging.info("Uploading to Signed URL: %s", gcs_path) + with open(local_path, "rb") as f: + data = f.read() + req = urllib.request.Request(gcs_path, data=data, method="PUT") + with urllib.request.urlopen(req) as response: + if response.status not in (200, 201): + raise RuntimeError(f"Upload failed with status {response.status}: {response.read()}") + return + parts = gcs_path.replace("gs://", "").split("/", 1) bucket_name = parts[0] blob_path = parts[1] diff --git a/kinetic/runner/remote_runner_test.py b/kinetic/runner/remote_runner_test.py index 7b9dd30f..595efcd6 100644 --- a/kinetic/runner/remote_runner_test.py +++ b/kinetic/runner/remote_runner_test.py @@ -6,6 +6,8 @@ import shutil import sys import tempfile +import urllib.request +import urllib.error import zipfile from unittest import mock from unittest.mock import MagicMock @@ -1093,5 +1095,96 @@ def fake_download(client, gcs_path, local_path): self.assertEqual(cm.exception.code, 1) +class TestHttpOperations(absltest.TestCase): + """Verify that download and upload helpers support Signed URLs.""" + + @mock.patch("kinetic.runner.remote_runner.urllib.request.urlretrieve") + def test_download_from_signed_url(self, mock_urlretrieve): + # If it's a signed URL (starts with http/https), it should use urlretrieve + _download_from_gcs( + client=None, + gcs_path="https://storage.googleapis.com/bucket/job/payload.pkl?sig=123", + local_path="/tmp/local.pkl" + ) + mock_urlretrieve.assert_called_once_with( + "https://storage.googleapis.com/bucket/job/payload.pkl?sig=123", + "/tmp/local.pkl" + ) + + @mock.patch("kinetic.runner.remote_runner.urllib.request.urlopen") + def test_upload_to_signed_url(self, mock_urlopen): + # Mock urlopen response + mock_response = mock.MagicMock() + mock_response.status = 200 + mock_urlopen.return_value.__enter__.return_value = mock_response + + # Create a dummy file to upload + with tempfile.NamedTemporaryFile(delete=False) as f: + f.write(b"test data") + temp_file_path = f.name + self.addCleanup(os.unlink, temp_file_path) + + _upload_to_gcs( + client=None, + local_path=temp_file_path, + gcs_path="https://storage.googleapis.com/bucket/job/result.pkl?sig=123" + ) + + # Verify urlopen was called with a PUT request containing the file data + mock_urlopen.assert_called_once() + req = mock_urlopen.call_args[0][0] + self.assertIsInstance(req, urllib.request.Request) + self.assertEqual(req.full_url, "https://storage.googleapis.com/bucket/job/result.pkl?sig=123") + self.assertEqual(req.method, "PUT") + self.assertEqual(req.data, b"test data") + + +class TestSentinels(absltest.TestCase): + """Verify sentinel upload and wait operations using Signed URLs.""" + + @mock.patch("kinetic.runner.remote_runner.urllib.request.urlopen") + def test_upload_sentinel(self, mock_urlopen): + mock_response = mock.MagicMock() + mock_response.status = 200 + mock_urlopen.return_value.__enter__.return_value = mock_response + + from kinetic.runner.remote_runner import _upload_sentinel + _upload_sentinel("https://signed-url/sentinel") + + mock_urlopen.assert_called_once() + req = mock_urlopen.call_args[0][0] + self.assertEqual(req.full_url, "https://signed-url/sentinel") + self.assertEqual(req.method, "PUT") + self.assertEqual(req.data, b"") + + @mock.patch("kinetic.runner.remote_runner.urllib.request.urlopen") + @mock.patch("kinetic.runner.remote_runner.time.sleep") # speed up test + def test_wait_for_sentinel_success(self, mock_sleep, mock_urlopen): + mock_response = mock.MagicMock() + mock_response.status = 200 + mock_urlopen.return_value.__enter__.return_value = mock_response + + from kinetic.runner.remote_runner import _wait_for_sentinel + # Should return immediately if 200 OK + _wait_for_sentinel("https://signed-url/sentinel", timeout=10) + mock_urlopen.assert_called_once() + + @mock.patch("kinetic.runner.remote_runner.urllib.request.urlopen") + @mock.patch("kinetic.runner.remote_runner.time.sleep") + def test_wait_for_sentinel_timeout(self, mock_sleep, mock_urlopen): + # Simulate 404 Not Found on every check + mock_urlopen.side_effect = urllib.error.HTTPError( + url="https://signed-url/sentinel", + code=404, + msg="Not Found", + hdrs=None, + fp=None + ) + + from kinetic.runner.remote_runner import _wait_for_sentinel + with self.assertRaises(RuntimeError): + _wait_for_sentinel("https://signed-url/sentinel", timeout=2) # low timeout for test + + if __name__ == "__main__": absltest.main() diff --git a/kinetic/utils/storage.py b/kinetic/utils/storage.py index e1a79712..75ffd265 100644 --- a/kinetic/utils/storage.py +++ b/kinetic/utils/storage.py @@ -2,6 +2,7 @@ from __future__ import annotations +import datetime import json import os import tempfile @@ -13,6 +14,9 @@ from google.cloud.storage import transfer_manager from google.cloud.storage.retry import DEFAULT_RETRY +import google.auth +from google.auth import impersonated_credentials + from kinetic.constants import get_default_project from kinetic.data import Data @@ -391,3 +395,99 @@ def _upload_directory( worker_type=transfer_manager.THREAD, raise_exception=True, ) + + +def generate_job_signed_urls( + bucket_name: str, + job_id: str, + project: str, + signer_sa_email: str, + has_requirements: bool = False, + debug: bool = False, +) -> dict[str, str]: + """Generate GCS Signed URLs for job isolation. + + Impersonates the signer service account to sign URLs for job artifacts, + preventing the worker pods from needing broad GCS access. + + Args: + bucket_name: GCS bucket name. + job_id: Unique job identifier. + project: GCP project ID. + signer_sa_email: GSA email to impersonate for signing. + has_requirements: Whether the job has a requirements.txt to download. + debug: Whether debug mode is enabled (generates debug sentinel URLs). + + Returns: + Dict mapping artifact names to their signed URLs. + """ + # 1. Get source credentials (user or default GSA) + source_credentials, _ = google.auth.default() + + # 2. Create impersonated credentials for the signer SA + target_scopes = ["https://www.googleapis.com/auth/devstorage.full_control"] + creds = impersonated_credentials.Credentials( + source_credentials=source_credentials, + target_principal=signer_sa_email, + target_scopes=target_scopes, + lifetime=3600, + ) + + # 3. Initialize storage client with impersonated credentials + client = storage.Client(project=project, credentials=creds) + bucket = client.bucket(bucket_name) + + # 4. Generate signed URLs + expiration = datetime.timedelta(days=7) + urls = {} + + # Standard artifacts + for name, key, method in [ + ("payload.pkl", "payload_download", "GET"), + ("context.zip", "context_download", "GET"), + ("result.pkl", "result_upload", "PUT"), + ]: + blob = bucket.blob(f"{job_id}/{name}") + urls[key] = blob.generate_signed_url( + version="v4", + expiration=expiration, + method=method, + service_account_email=signer_sa_email, + ) + + # Optional requirements + if has_requirements: + blob = bucket.blob(f"{job_id}/requirements.txt") + urls["requirements_download"] = blob.generate_signed_url( + version="v4", + expiration=expiration, + method="GET", + service_account_email=signer_sa_email, + ) + + # Optional debug sentinels + if debug: + # .debug_ready (PUT) + blob = bucket.blob(f"{job_id}/.debug_ready") + urls["debug_ready_upload"] = blob.generate_signed_url( + version="v4", + expiration=expiration, + method="PUT", + service_account_email=signer_sa_email, + ) + # .leader_ready (PUT for leader, GET for workers) + blob = bucket.blob(f"{job_id}/.leader_ready") + urls["leader_ready_upload"] = blob.generate_signed_url( + version="v4", + expiration=expiration, + method="PUT", + service_account_email=signer_sa_email, + ) + urls["leader_ready_download"] = blob.generate_signed_url( + version="v4", + expiration=expiration, + method="GET", + service_account_email=signer_sa_email, + ) + + return urls diff --git a/kinetic/utils/storage_test.py b/kinetic/utils/storage_test.py index 0cfd9ae3..a416abcb 100644 --- a/kinetic/utils/storage_test.py +++ b/kinetic/utils/storage_test.py @@ -389,5 +389,92 @@ def test_empty_directory_is_noop(self): self.mock_upload.assert_not_called() +class TestGenerateSignedUrls(_GcsTestBase): + """Verify generation of GCS Signed URLs for job isolation.""" + + @mock.patch("kinetic.utils.storage.storage.Client") + @mock.patch("kinetic.utils.storage.impersonated_credentials.Credentials") + @mock.patch("kinetic.utils.storage.google.auth.default") + def test_generates_standard_urls(self, mock_auth_default, mock_impersonated_creds, mock_client_class): + # Mock default credentials + mock_user_creds = mock.MagicMock() + mock_auth_default.return_value = (mock_user_creds, "test-project") + + # Mock impersonated credentials + mock_imp_creds = mock_impersonated_creds.return_value + + # Mock GCS blob generate_signed_url + mock_bucket = mock_client_class.return_value.bucket.return_value + mock_blob = mock_bucket.blob.return_value + mock_blob.generate_signed_url.return_value = "https://signed-url" + + urls = storage_module.generate_job_signed_urls( + bucket_name="my-bucket", + job_id="job-123", + project="test-project", + signer_sa_email="signer@test.iam.gserviceaccount.com", + has_requirements=False, + debug=False + ) + + # Verify impersonation was set up + mock_impersonated_creds.assert_called_once_with( + source_credentials=mock_user_creds, + target_principal="signer@test.iam.gserviceaccount.com", + target_scopes=["https://www.googleapis.com/auth/devstorage.full_control"], + lifetime=3600 + ) + + # Verify storage.Client was initialized with impersonated credentials + mock_client_class.assert_called_with(project="test-project", credentials=mock_imp_creds) + + # Verify standard URLs are in the result + self.assertEqual(urls["payload_download"], "https://signed-url") + self.assertEqual(urls["context_download"], "https://signed-url") + self.assertEqual(urls["result_upload"], "https://signed-url") + self.assertNotIn("requirements_download", urls) + self.assertNotIn("debug_ready_upload", urls) + + # Verify generate_signed_url calls + self.assertEqual(mock_blob.generate_signed_url.call_count, 3) + mock_blob.generate_signed_url.assert_any_call( + version="v4", + expiration=mock.ANY, + method="GET", + service_account_email="signer@test.iam.gserviceaccount.com" + ) + mock_blob.generate_signed_url.assert_any_call( + version="v4", + expiration=mock.ANY, + method="PUT", + service_account_email="signer@test.iam.gserviceaccount.com" + ) + + @mock.patch("kinetic.utils.storage.impersonated_credentials.Credentials") + @mock.patch("kinetic.utils.storage.google.auth.default") + def test_generates_optional_urls(self, mock_auth_default, mock_impersonated_creds): + mock_user_creds = mock.MagicMock() + mock_auth_default.return_value = (mock_user_creds, "test-project") + + mock_bucket = self.mock_gcs.bucket.return_value + mock_blob = mock_bucket.blob.return_value + mock_blob.generate_signed_url.return_value = "https://signed-url" + + urls = storage_module.generate_job_signed_urls( + bucket_name="my-bucket", + job_id="job-123", + project="test-project", + signer_sa_email="signer@test.iam.gserviceaccount.com", + has_requirements=True, + debug=True + ) + + self.assertIn("requirements_download", urls) + self.assertIn("debug_ready_upload", urls) + self.assertIn("leader_ready_upload", urls) + self.assertIn("leader_ready_download", urls) + self.assertEqual(mock_blob.generate_signed_url.call_count, 7) # 3 std + 1 req + 3 debug + + if __name__ == "__main__": absltest.main()