From c0bfbf8d128ab83c3dd19bb41c30e0ddd4dedbdc Mon Sep 17 00:00:00 2001 From: Keita Watanabe Date: Sun, 29 Mar 2026 00:08:11 +0000 Subject: [PATCH 1/2] feat: support DefaultDataCredential for workflow_log and workflow_data MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Widen LogConfig, DataConfig, and related function signatures from StaticDataCredential to DataCredential (the union that includes DefaultDataCredential). This allows workflow log/data storage to use ambient credentials (IRSA, Pod Identity, instance metadata) instead of requiring static IAM access keys. Also guard the Go sidecar's env var override so empty static keys don't clobber the SDK's ambient credential chain. Python changes (postgres.py, task.py): type annotation only — all downstream code already uses DataCredential-compatible interfaces (to_decrypted_dict, storage.Client.create). Go change (data.go): skip os.Setenv when AccessKeyId is empty. Fixes NVIDIA/OSMO#749 --- src/runtime/pkg/data/data.go | 10 ++++++++-- src/utils/connectors/postgres.py | 6 +++--- src/utils/job/task.py | 4 ++-- 3 files changed, 13 insertions(+), 7 deletions(-) diff --git a/src/runtime/pkg/data/data.go b/src/runtime/pkg/data/data.go index 22a6f0e95..53f478f0c 100644 --- a/src/runtime/pkg/data/data.go +++ b/src/runtime/pkg/data/data.go @@ -414,8 +414,14 @@ func MountURL(downloadType string, credentialInfo ConfigInfo, urlPath string, osmoChan <- fmt.Sprintf("Missing data credential for %s.", storageBackend.GetProfile()) return isEmpty } - os.Setenv("AWS_ACCESS_KEY_ID", dataCredential.AccessKeyId) - os.Setenv("AWS_SECRET_ACCESS_KEY", dataCredential.AccessKey) + // Only set static key env vars when keys are provided. + // When using DefaultDataCredential (ambient credentials via Pod Identity, + // IRSA, etc.), keys are empty — setting empty env vars would clobber the + // SDK's default credential chain. + if dataCredential.AccessKeyId != "" { + os.Setenv("AWS_ACCESS_KEY_ID", dataCredential.AccessKeyId) + os.Setenv("AWS_SECRET_ACCESS_KEY", dataCredential.AccessKey) + } var commandArgs []string diff --git a/src/utils/connectors/postgres.py b/src/utils/connectors/postgres.py index d12d68d6f..a1bcb99c6 100644 --- a/src/utils/connectors/postgres.py +++ b/src/utils/connectors/postgres.py @@ -1480,7 +1480,7 @@ def get_data_cred(self, user: str, profile: str) -> credentials.StaticDataCreden return None - def get_all_data_creds(self, user: str) -> Dict[str, credentials.StaticDataCredential]: + def get_all_data_creds(self, user: str) -> Dict[str, credentials.DataCredential]: """ Fetch all data credentials for user. """ select_data_cmd = PostgresSelectCommand( table='credential', @@ -2557,7 +2557,7 @@ def construct_path(endpoint: str, bucket: str, path: str): class LogConfig(ExtraArgBaseModel): """ Config for storing information about data. """ - credential: credentials.StaticDataCredential | None = None + credential: credentials.DataCredential | None = None class WorkflowInfo(ExtraArgBaseModel): @@ -2574,7 +2574,7 @@ def validate_name(self, name: str): class DataConfig(ExtraArgBaseModel): """ Config for storing information about data. """ - credential: credentials.StaticDataCredential | None = None + credential: credentials.DataCredential | None = None base_url: str = '' # Timeout in mins for osmo-ctrl to retry connecting to the OSMO service until exiting the task diff --git a/src/utils/job/task.py b/src/utils/job/task.py index 2bbafe9c7..0469d8ac0 100644 --- a/src/utils/job/task.py +++ b/src/utils/job/task.py @@ -102,7 +102,7 @@ def create_login_dict(user: str, def create_config_dict( - data_info: dict[str, credentials.StaticDataCredential], + data_info: dict[str, credentials.DataCredential], ) -> dict: ''' Creates the config dict where the input should be a dict containing key values like: @@ -2339,7 +2339,7 @@ def convert_to_pod_spec( service_config: connectors.ServiceConfig | None = None, dataset_config: connectors.DatasetConfig | None = None, pool_info: connectors.Pool | None = None, - data_endpoints: Dict[str, credentials.StaticDataCredential] | None = None, + data_endpoints: Dict[str, credentials.DataCredential] | None = None, skip_refresh_token: bool = False, auth_token: str | None = None, ) -> Tuple[Dict, Dict[str, kb_objects.FileMount], Optional[Tuple[str, str]]]: From 7e84b51c64a78372b9977c892b554c3bd811eab8 Mon Sep 17 00:00:00 2001 From: Keita Watanabe Date: Sun, 29 Mar 2026 00:14:25 +0000 Subject: [PATCH 2/2] fix: revert get_all_data_creds return type and set AWS_REGION in sidecar - Revert get_all_data_creds return type to StaticDataCredential (it constructs StaticDataCredential explicitly, not DataCredential) - Set AWS_REGION env var from dataCredential.Region when non-empty, so ambient credentials (IRSA/Pod Identity) can locate the correct S3 endpoint --- src/runtime/pkg/data/data.go | 3 +++ src/utils/connectors/postgres.py | 2 +- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/src/runtime/pkg/data/data.go b/src/runtime/pkg/data/data.go index 53f478f0c..3d68a5bc0 100644 --- a/src/runtime/pkg/data/data.go +++ b/src/runtime/pkg/data/data.go @@ -422,6 +422,9 @@ func MountURL(downloadType string, credentialInfo ConfigInfo, urlPath string, os.Setenv("AWS_ACCESS_KEY_ID", dataCredential.AccessKeyId) os.Setenv("AWS_SECRET_ACCESS_KEY", dataCredential.AccessKey) } + if dataCredential.Region != "" { + os.Setenv("AWS_REGION", dataCredential.Region) + } var commandArgs []string diff --git a/src/utils/connectors/postgres.py b/src/utils/connectors/postgres.py index a1bcb99c6..56f7ff7ae 100644 --- a/src/utils/connectors/postgres.py +++ b/src/utils/connectors/postgres.py @@ -1480,7 +1480,7 @@ def get_data_cred(self, user: str, profile: str) -> credentials.StaticDataCreden return None - def get_all_data_creds(self, user: str) -> Dict[str, credentials.DataCredential]: + def get_all_data_creds(self, user: str) -> Dict[str, credentials.StaticDataCredential]: """ Fetch all data credentials for user. """ select_data_cmd = PostgresSelectCommand( table='credential',