outerbounds 0.3.88__py3-none-any.whl → 0.3.90__py3-none-any.whl

outerbounds/command_groups/local_setup_cli.py

@@ -640,7 +640,6 @@ class ConfigurationWriter:
         if config_type == "inline":
             if "OBP_PERIMETER" in self.decoded_config:
                 self.selected_perimeter = self.decoded_config["OBP_PERIMETER"]
-
             if "OBP_METAFLOW_CONFIG_URL" in self.decoded_config:
                 self.decoded_config = {
                     "OBP_METAFLOW_CONFIG_URL": self.decoded_config[
@@ -701,8 +700,6 @@ class ConfigurationWriter:
         with open(config_path, "w") as fd:
             json.dump(self.existing, fd, indent=4)
 
-        # Every time a config is initialized, we should also reset the corresponding ob_config[_profile].json
-        remote_config = metaflowconfig.init_config(self.out_dir, self.profile)
         if self.selected_perimeter and "OBP_METAFLOW_CONFIG_URL" in self.decoded_config:
             with open(self.ob_config_path, "w") as fd:
                 ob_config_dict = {

outerbounds/command_groups/perimeters_cli.py

@@ -1,21 +1,15 @@
-import base64
-import hashlib
 import json
 import os
-import re
-import subprocess
 import sys
-import zlib
-from base64 import b64decode, b64encode
-from importlib.machinery import PathFinder
+from io import StringIO
 from os import path
-from pathlib import Path
-from typing import Any, Callable, Dict, List
+from typing import Any, Dict
 from outerbounds._vendor import click
 import requests
-from requests.exceptions import HTTPError
+import configparser
 
-from ..utils import kubeconfig, metaflowconfig
+from ..utils import metaflowconfig
+from ..utils.utils import safe_write_to_disk
 from ..utils.schema import (
     CommandStatus,
     OuterboundsCommandResponse,
@@ -80,7 +74,7 @@ def switch(config_dir=None, profile=None, output="", id=None, force=False):
         id, perimeters, output, switch_perimeter_response, switch_perimeter_step
     )
 
-    path_to_config = get_ob_config_file_path(config_dir, profile)
+    path_to_config = metaflowconfig.get_ob_config_file_path(config_dir, profile)
 
     import fcntl
 
@@ -124,9 +118,31 @@ def switch(config_dir=None, profile=None, output="", id=None, force=False):
         )
 
     switch_perimeter_response.add_step(switch_perimeter_step)
+
+    ensure_cloud_creds_step = CommandStatus(
+        "EnsureCloudCredentials",
+        OuterboundsCommandStatus.OK,
+        "Cloud credentials were successfully updated.",
+    )
+
+    try:
+        ensure_cloud_credentials_for_shell(config_dir, profile)
+    except:
+        click.secho(
+            "Failed to update cloud credentials.",
+            fg="red",
+            err=True,
+        )
+        ensure_cloud_creds_step.update(
+            status=OuterboundsCommandStatus.FAIL,
+            reason="Failed to update cloud credentials.",
+            mitigation="",
+        )
+
+    switch_perimeter_response.add_step(ensure_cloud_creds_step)
+
     if output == "json":
         click.echo(json.dumps(switch_perimeter_response.as_dict(), indent=4))
-        return
 
 
 @perimeter.command(help="Show current perimeter")
@@ -275,6 +291,58 @@ def list(config_dir=None, profile=None, output=""):
         click.echo(json.dumps(list_perimeters_response.as_dict(), indent=4))
 
 
+@perimeter.command(
+    help="Ensure credentials for cloud are synced with perimeter", hidden=True
+)
+@click.option(
+    "-d",
+    "--config-dir",
+    default=path.expanduser(os.environ.get("METAFLOW_HOME", "~/.metaflowconfig")),
+    help="Path to Metaflow configuration directory",
+    show_default=True,
+)
+@click.option(
+    "-p",
+    "--profile",
+    default=os.environ.get("METAFLOW_PROFILE", ""),
+    help="The named metaflow profile in which your workstation exists",
+)
+@click.option(
+    "-o",
+    "--output",
+    default="",
+    help="Show output in the specified format.",
+    type=click.Choice(["json", ""]),
+)
+def ensure_cloud_creds(config_dir=None, profile=None, output=""):
+    ensure_cloud_creds_step = CommandStatus(
+        "EnsureCloudCredentials",
+        OuterboundsCommandStatus.OK,
+        "Cloud credentials were successfully updated.",
+    )
+
+    ensure_cloud_creds_response = OuterboundsCommandResponse()
+
+    try:
+        ensure_cloud_credentials_for_shell(config_dir, profile)
+        click.secho("Cloud credentials updated successfully.", fg="green", err=True)
+    except:
+        click.secho(
+            "Failed to update cloud credentials.",
+            fg="red",
+            err=True,
+        )
+        ensure_cloud_creds_step.update(
+            status=OuterboundsCommandStatus.FAIL,
+            reason="Failed to update cloud credentials.",
+            mitigation="",
+        )
+
+    ensure_cloud_creds_response.add_step(ensure_cloud_creds_step)
+    if output == "json":
+        click.echo(json.dumps(ensure_cloud_creds_response.as_dict(), indent=4))
+
+
 def get_list_perimeters_api_response(config_dir, profile):
     metaflow_token = metaflowconfig.get_metaflow_token_from_config(config_dir, profile)
     api_url = metaflowconfig.get_sanitized_url_from_config(
@@ -288,15 +356,6 @@ def get_list_perimeters_api_response(config_dir, profile):
     return perimeters_response.json()["perimeters"]
 
 
-def get_ob_config_file_path(config_dir: str, profile: str) -> str:
-    # If OBP_CONFIG_DIR is set, use that, otherwise use METAFLOW_HOME
-    # If neither are set, use ~/.metaflowconfig
-    obp_config_dir = path.expanduser(os.environ.get("OBP_CONFIG_DIR", config_dir))
-
-    ob_config_filename = f"ob_config_{profile}.json" if profile else "ob_config.json"
-    return os.path.expanduser(os.path.join(obp_config_dir, ob_config_filename))
-
-
 def get_perimeters_from_api_or_fail_command(
     config_dir: str,
     profile: str,
@@ -331,7 +390,7 @@ def get_ob_config_or_fail_command(
     command_response: OuterboundsCommandResponse,
     command_step: CommandStatus,
 ) -> Dict[str, str]:
-    path_to_config = get_ob_config_file_path(config_dir, profile)
+    path_to_config = metaflowconfig.get_ob_config_file_path(config_dir, profile)
 
     if not os.path.exists(path_to_config):
         click.secho(
@@ -371,6 +430,23 @@ def get_ob_config_or_fail_command(
     return ob_config_dict
 
 
+def ensure_cloud_credentials_for_shell(config_dir, profile):
+    if "WORKSTATION_ID" not in os.environ:
+        # Naive check to see if we're running in workstation. No need to ensure anything
+        # if this is not a workstation.
+        return
+
+    mf_config = metaflowconfig.init_config(config_dir, profile)
+
+    # Currently we only support GCP. TODO: utkarsh to add support for AWS and Azure
+    if "METAFLOW_DEFAULT_GCP_CLIENT_PROVIDER" in mf_config:
+        # This is a GCP deployment.
+        ensure_gcp_cloud_creds(config_dir, profile)
+    elif "METAFLOW_DEFAULT_AWS_CLIENT_PROVIDER" in mf_config:
+        # This is an AWS deployment.
+        ensure_aws_cloud_creds(config_dir, profile)
+
+
 def confirm_user_has_access_to_perimeter_or_fail(
     perimeter_id: str,
     perimeters: Dict[str, Any],
@@ -395,4 +471,93 @@ def confirm_user_has_access_to_perimeter_or_fail(
         sys.exit(1)
 
 
+def ensure_gcp_cloud_creds(config_dir, profile):
+    token_info = get_gcp_auth_credentials(config_dir, profile)
+    auth_url = metaflowconfig.get_sanitized_url_from_config(
+        config_dir, profile, "OBP_AUTH_SERVER"
+    )
+    metaflow_token = metaflowconfig.get_metaflow_token_from_config(config_dir, profile)
+
+    try:
+        # GOOGLE_APPLICATION_CREDENTIALS is a well known gcloud environment variable
+        credentials_file_loc = os.environ["GOOGLE_APPLICATION_CREDENTIALS"]
+    except KeyError:
+        # This is most likely an old workstation when these params weren't set. Do nothing.
+        # Alternatively, user might have deliberately unset it to use their own auth.
+        return
+
+    credentials_json = {
+        "type": "external_account",
+        "audience": f"//iam.googleapis.com/projects/{token_info['gcpProjectNumber']}/locations/global/workloadIdentityPools/{token_info['gcpWorkloadIdentityPool']}/providers/{token_info['gcpWorkloadIdentityPoolProvider']}",
+        "subject_token_type": "urn:ietf:params:oauth:token-type:jwt",
+        "token_url": "https://sts.googleapis.com/v1/token",
+        "service_account_impersonation_url": f"https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/{token_info['gcpServiceAccountEmail']}:generateAccessToken",
+        "credential_source": {
+            "url": f"{auth_url}/generate/gcp",
+            "headers": {"x-api-key": metaflow_token},
+            "format": {"type": "json", "subject_token_field_name": "token"},
+        },
+    }
+
+    safe_write_to_disk(credentials_file_loc, json.dumps(credentials_json))
+
+
+def ensure_aws_cloud_creds(config_dir, profile):
+    token_info = get_aws_auth_credentials(config_dir, profile)
+
+    try:
+        token_file_loc = os.environ["OBP_AWS_WEB_IDENTITY_TOKEN_FILE"]
+
+        # AWS_CONFIG_FILE is a well known aws cli environment variable
+        config_file_loc = os.environ["AWS_CONFIG_FILE"]
+    except KeyError:
+        # This is most likely an old workstation when these params weren't set. Do nothing.
+        # Alternatively, user might have deliberately unset it to use their own auth.
+        return
+
+    aws_config = configparser.ConfigParser()
+    aws_config.read(config_file_loc)
+
+    profile_name = "profile outerbounds"
+    if profile_name not in aws_config:
+        aws_config[profile_name] = {}
+
+    aws_config[profile_name]["role_arn"] = token_info["role_arn"]
+    aws_config[profile_name]["web_identity_token_file"] = token_file_loc
+
+    aws_config_string = StringIO()
+    aws_config.write(aws_config_string)
+
+    safe_write_to_disk(token_file_loc, token_info["token"])
+    safe_write_to_disk(config_file_loc, aws_config_string.getvalue())
+
+
+def get_aws_auth_credentials(config_dir, profile):  # pragma: no cover
+    token = metaflowconfig.get_metaflow_token_from_config(config_dir, profile)
+    auth_server_url = metaflowconfig.get_sanitized_url_from_config(
+        config_dir, profile, "OBP_AUTH_SERVER"
+    )
+
+    response = requests.get(
+        "{}/generate/aws".format(auth_server_url), headers={"x-api-key": token}
+    )
+    response.raise_for_status()
+
+    return response.json()
+
+
+def get_gcp_auth_credentials(config_dir, profile):
+    token = metaflowconfig.get_metaflow_token_from_config(config_dir, profile)
+    auth_server_url = metaflowconfig.get_sanitized_url_from_config(
+        config_dir, profile, "OBP_AUTH_SERVER"
+    )
+
+    response = requests.get(
+        "{}/generate/gcp".format(auth_server_url), headers={"x-api-key": token}
+    )
+    response.raise_for_status()
+
+    return response.json()
+
+
 cli.add_command(perimeter, name="perimeter")

outerbounds/utils/metaflowconfig.py

@@ -3,21 +3,43 @@ import json
 import os
 import requests
 from os import path
-from typing import Dict
+from typing import Dict, Union
 import sys
 
+"""
+key: perimeter specific URL to fetch the remote metaflow config from
+value: the remote metaflow config
+"""
+CACHED_REMOTE_METAFLOW_CONFIG: Dict[str, Dict[str, str]] = {}
+
+
+CURRENT_PERIMETER_KEY = "OB_CURRENT_PERIMETER"
+CURRENT_PERIMETER_URL = "OB_CURRENT_PERIMETER_MF_CONFIG_URL"
+CURRENT_PERIMETER_URL_LEGACY_KEY = (
+    "OB_CURRENT_PERIMETER_URL"  # For backwards compatibility with workstations.
+)
+
 
 def init_config(config_dir, profile) -> Dict[str, str]:
+    global CACHED_REMOTE_METAFLOW_CONFIG
     config = read_metaflow_config_from_filesystem(config_dir, profile)
 
-    # This is new remote-metaflow config; fetch it from the URL
-    if "OBP_METAFLOW_CONFIG_URL" in config:
-        remote_config = init_config_from_url(
-            config_dir, profile, config["OBP_METAFLOW_CONFIG_URL"]
-        )
-        remote_config["OBP_METAFLOW_CONFIG_URL"] = config["OBP_METAFLOW_CONFIG_URL"]
+    # Either user has an ob_config.json file with the perimeter URL
+    # or the default config on the filesystem has the config URL in it.
+    perimeter_specifc_url = get_perimeter_config_url_if_set_in_ob_config(
+        config_dir, profile
+    ) or config.get("OBP_METAFLOW_CONFIG_URL", "")
+
+    if perimeter_specifc_url != "":
+        if perimeter_specifc_url in CACHED_REMOTE_METAFLOW_CONFIG:
+            return CACHED_REMOTE_METAFLOW_CONFIG[perimeter_specifc_url]
+
+        remote_config = init_config_from_url(config_dir, profile, perimeter_specifc_url)
+        remote_config["OBP_METAFLOW_CONFIG_URL"] = perimeter_specifc_url
+
+        CACHED_REMOTE_METAFLOW_CONFIG[perimeter_specifc_url] = remote_config
         return remote_config
-    # Legacy config, use from filesystem
+
     return config
 
 
@@ -101,3 +123,40 @@ def get_remote_metaflow_config_for_perimeter(
             fg="red",
         )
         sys.exit(1)
+
+
+def get_ob_config_file_path(config_dir: str, profile: str) -> str:
+    # If OBP_CONFIG_DIR is set, use that, otherwise use METAFLOW_HOME
+    # If neither are set, use ~/.metaflowconfig
+    obp_config_dir = path.expanduser(os.environ.get("OBP_CONFIG_DIR", config_dir))
+
+    ob_config_filename = f"ob_config_{profile}.json" if profile else "ob_config.json"
+    return os.path.expanduser(os.path.join(obp_config_dir, ob_config_filename))
+
+
+def get_perimeter_config_url_if_set_in_ob_config(
+    config_dir: str, profile: str
+) -> Union[str, None]:
+    file_path = get_ob_config_file_path(config_dir, profile)
+
+    if os.path.exists(file_path):
+        with open(file_path, "r") as f:
+            ob_config = json.loads(f.read())
+
+        if CURRENT_PERIMETER_URL in ob_config:
+            return ob_config[CURRENT_PERIMETER_URL]
+        elif CURRENT_PERIMETER_URL_LEGACY_KEY in ob_config:
+            return ob_config[CURRENT_PERIMETER_URL_LEGACY_KEY]
+        else:
+            raise ValueError(
+                "{} does not contain the key {}".format(
+                    file_path, CURRENT_PERIMETER_KEY
+                )
+            )
+    elif "OBP_CONFIG_DIR" in os.environ:
+        raise FileNotFoundError(
+            "Environment variable OBP_CONFIG_DIR is set to {} but this directory does not contain an ob_config.json file.".format(
+                os.environ["OBP_CONFIG_DIR"]
+            )
+        )
+    return None

outerbounds/utils/utils.py

@@ -0,0 +1,19 @@
+import tempfile
+import os
+
+"""
+Writes the given data to file_loc. Ensures that the directory exists
+and uses a temporary file to ensure that the file is written atomically.
+"""
+
+
+def safe_write_to_disk(file_loc, data):
+    # Ensure the directory exists
+    os.makedirs(os.path.dirname(file_loc), exist_ok=True)
+
+    with tempfile.NamedTemporaryFile(
+        "w", dir=os.path.dirname(file_loc), delete=False
+    ) as f:
+        f.write(data)
+        tmp_file = f.name
+    os.rename(tmp_file, file_loc)

{outerbounds-0.3.88.dist-info → outerbounds-0.3.90.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: outerbounds
-Version: 0.3.88
+Version: 0.3.90
 Summary: More Data Science, Less Administration
 License: Proprietary
 Keywords: data science,machine learning,MLOps

{outerbounds-0.3.88.dist-info → outerbounds-0.3.90.dist-info}/RECORD

@@ -42,15 +42,16 @@ outerbounds/_vendor/yaml/tokens.py,sha256=JBSu38wihGr4l73JwbfMA7Ks1-X84g8-NskTz7
 outerbounds/cli_main.py,sha256=e9UMnPysmc7gbrimq2I4KfltggyU7pw59Cn9aEguVcU,74
 outerbounds/command_groups/__init__.py,sha256=QPWtj5wDRTINDxVUL7XPqG3HoxHNvYOg08EnuSZB2Hc,21
 outerbounds/command_groups/cli.py,sha256=sorDdQvmTPqIwfvgtuNLILelimXu5CknFnWQFsYFGHs,286
-outerbounds/command_groups/local_setup_cli.py,sha256=somQMeLgRSK8BAje2rN6LeY-lszXmwBpNLvDCk293h8,36554
-outerbounds/command_groups/perimeters_cli.py,sha256=vU1LykO9T2Xj4Fn8hyhaCUR4q454q1aVoU4XhhrgAS4,12781
+outerbounds/command_groups/local_setup_cli.py,sha256=tuuqJRXQ_guEwOuQSIf9wkUU0yg8yAs31myGViAK15s,36364
+outerbounds/command_groups/perimeters_cli.py,sha256=mrJfFIRYFOjuiz-9h4OKg2JT8Utmbs72z6wvPzDss3s,18685
 outerbounds/command_groups/workstations_cli.py,sha256=V5Jbj1cVb4IRllI7fOgNgL6OekRpuFDv6CEhDb4xC6w,22016
 outerbounds/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 outerbounds/utils/kubeconfig.py,sha256=yvcyRXGR4AhQuqUDqmbGxEOHw5ixMFV0AZIDg1LI_Qo,7981
-outerbounds/utils/metaflowconfig.py,sha256=Ja1eufYm4UANRNpk7qlxg_UUudToDOVCjaWg88N-xhQ,3538
+outerbounds/utils/metaflowconfig.py,sha256=l2vJbgPkLISU-XPGZFaC8ZKmYFyJemlD6bwB-EKUsAw,5770
 outerbounds/utils/schema.py,sha256=cNlgjmteLPbDzSEUSQDsq8txdhMGyezSmM83jU3aa0w,2329
+outerbounds/utils/utils.py,sha256=4Z8cszNob_8kDYCLNTrP-wWads_S_MdL3Uj3ju4mEsk,501
 outerbounds/vendor.py,sha256=gRLRJNXtZBeUpPEog0LOeIsl6GosaFFbCxUvR4bW6IQ,5093
-outerbounds-0.3.88.dist-info/entry_points.txt,sha256=7ye0281PKlvqxu15rjw60zKg2pMsXI49_A8BmGqIqBw,47
-outerbounds-0.3.88.dist-info/WHEEL,sha256=vVCvjcmxuUltf8cYhJ0sJMRDLr1XsPuxEId8YDzbyCY,88
-outerbounds-0.3.88.dist-info/METADATA,sha256=Pg9IyVEEVQpR91fgtQL43o0fBVwNzGUBJ0GSprAT4DY,1632
-outerbounds-0.3.88.dist-info/RECORD,,
+outerbounds-0.3.90.dist-info/entry_points.txt,sha256=7ye0281PKlvqxu15rjw60zKg2pMsXI49_A8BmGqIqBw,47
+outerbounds-0.3.90.dist-info/WHEEL,sha256=vVCvjcmxuUltf8cYhJ0sJMRDLr1XsPuxEId8YDzbyCY,88
+outerbounds-0.3.90.dist-info/METADATA,sha256=an69uvqpICHuiBfJVZJJrZU474ehFkjnnilY1jRLYFg,1632
+outerbounds-0.3.90.dist-info/RECORD,,