metaflow 2.11.15__py2.py3-none-any.whl → 2.11.16__py2.py3-none-any.whl

metaflow/plugins/gcp/gcp_secret_manager_secrets_provider.py

@@ -0,0 +1,169 @@
+import base64
+import json
+from json import JSONDecodeError
+
+
+from metaflow.exception import MetaflowException
+from metaflow.plugins.secrets import SecretsProvider
+import re
+from metaflow.plugins.gcp.gs_storage_client_factory import get_credentials
+from metaflow.metaflow_config import GCP_SECRET_MANAGER_PREFIX
+
+
+class MetaflowGcpSecretsManagerBadResponse(MetaflowException):
+    """Raised when the response from GCP Secrets Manager is not valid in some way"""
+
+
+class MetaflowGcpSecretsManagerDuplicateKey(MetaflowException):
+    """Raised when the response from GCP Secrets Manager contains duplicate keys"""
+
+
+class MetaflowGcpSecretsManagerJSONParseError(MetaflowException):
+    """Raised when the SecretString response from GCP Secrets Manager is not valid JSON"""
+
+
+class MetaflowGcpSecretsManagerNotJSONObject(MetaflowException):
+    """Raised when the SecretString response from GCP Secrets Manager is not valid JSON dictionary"""
+
+
+def _sanitize_key_as_env_var(key):
+    """
+    Sanitize a key as an environment variable name.
+    This is purely a convenience trade-off to cover common cases well, vs. introducing
+    ambiguities (e.g. did the final '_' come from '.', or '-' or is original?).
+
+    1/27/2023(jackie):
+
+    We start with few rules and should *sparingly* add more over time.
+    Also, it's TBD whether all possible providers will share the same sanitization logic.
+    Therefore we will keep this function private for now
+    """
+    return key.replace("-", "_").replace(".", "_").replace("/", "_")
+
+
+class GcpSecretManagerSecretsProvider(SecretsProvider):
+    TYPE = "gcp-secret-manager"
+
+    def get_secret_as_dict(self, secret_id, options={}, role=None):
+        """
+        Reads a secret from GCP Secrets Manager and returns it as a dictionary of environment variables.
+
+        If the secret contains a string payload ("SecretString"):
+        - if the `json` option is True:
+            Secret will be parsed as a JSON. If successfully parsed, AND the JSON contains a
+            top-level object, each entry K/V in the object will also be converted to an entry in the result. V will
+            always be casted to a string (if not already a string).
+        - If `json` option is False (default):
+            Will be returned as a single entry in the result, with the key being the last part after / in secret_id.
+
+        On GCP Secrets Manager, the secret payload is a binary blob. However, by default we interpret it as UTF8 encoded
+        string. To disable this, set the `binary` option to True, the binary will be base64 encoded in the result.
+
+        All keys in the result are sanitized to be more valid environment variable names. This is done on a best effort
+        basis. Further validation is expected to be done by the invoking @secrets decorator itself.
+
+        :param secret_id: GCP Secrets Manager secret ID
+        :param options: unused
+        :return: dict of environment variables. All keys and values are strings.
+        """
+        from google.cloud.secretmanager_v1.services.secret_manager_service import (
+            SecretManagerServiceClient,
+        )
+        from google.cloud.secretmanager_v1.services.secret_manager_service.transports import (
+            SecretManagerServiceTransport,
+        )
+
+        # Full secret id looks like projects/1234567890/secrets/mysecret/versions/latest
+        #
+        # We allow these forms of secret_id:
+        #
+        # 1. Full path like projects/1234567890/secrets/mysecret/versions/latest
+        #    This is what you'd specify if you used to GCP SDK.
+        #
+        # 2. Full path but without the version like projects/1234567890/secrets/mysecret.
+        #    This is what you see in the GCP console, makes it easier to copy & paste.
+        #
+        # 3. Simple string like mysecret
+        #
+        # 4. Simple string with /versions/<version> suffix like mysecret/versions/1
+
+        # The latter two forms require METAFLOW_GCP_SECRET_MANAGER_PREFIX to be set.
+
+        match_full = re.match(
+            r"^projects/\d+/secrets/([\w\-]+)(/versions/([\w\-]+))?$", secret_id
+        )
+        match_partial = re.match(r"^([\w\-]+)(/versions/[\w\-]+)?$", secret_id)
+        if match_full:
+            # Full path
+            env_var_name = match_full.group(1)
+            if match_full.group(3):
+                # With version specified
+                full_secret_name = secret_id
+            else:
+                # No version specified, use latest
+                full_secret_name = secret_id + "/versions/latest"
+        elif match_partial:
+            # Partial path, possibly with /versions/<version> suffix
+            env_var_name = secret_id
+            if not GCP_SECRET_MANAGER_PREFIX:
+                raise ValueError(
+                    "Cannot use simple secret_id without setting METAFLOW_GCP_SECRET_MANAGER_PREFIX. %s"
+                    % GCP_SECRET_MANAGER_PREFIX
+                )
+            if match_partial.group(2):
+                # With version specified
+                full_secret_name = "%s%s" % (GCP_SECRET_MANAGER_PREFIX, secret_id)
+                env_var_name = match_partial.group(1)
+            else:
+                # No version specified, use latest
+                full_secret_name = "%s%s/versions/latest" % (
+                    GCP_SECRET_MANAGER_PREFIX,
+                    secret_id,
+                )
+        else:
+            raise ValueError(
+                "Invalid secret_id: %s. Must be either a full path or a simple string."
+                % secret_id
+            )
+
+        result = {}
+
+        def _sanitize_and_add_entry_to_result(k, v):
+            # Two jobs - sanitize, and check for dupes
+            sanitized_k = _sanitize_key_as_env_var(k)
+            if sanitized_k in result:
+                raise MetaflowGcpSecretsManagerDuplicateKey(
+                    "Duplicate key in secret: '%s' (sanitizes to '%s')"
+                    % (k, sanitized_k)
+                )
+            result[sanitized_k] = v
+
+        credentials, _ = get_credentials(
+            scopes=SecretManagerServiceTransport.AUTH_SCOPES
+        )
+        client = SecretManagerServiceClient(credentials=credentials)
+        response = client.access_secret_version(request={"name": full_secret_name})
+        payload_str = response.payload.data.decode("UTF-8")
+        if options.get("json", False):
+            obj = json.loads(payload_str)
+            if type(obj) == dict:
+                for k, v in obj.items():
+                    # We try to make it work here - cast to string always
+                    _sanitize_and_add_entry_to_result(k, str(v))
+            else:
+                raise MetaflowGcpSecretsManagerNotJSONObject(
+                    "Secret string is a JSON, but not an object (dict-like) - actual type %s."
+                    % type(obj)
+                )
+        else:
+            if options.get("env_var_name"):
+                env_var_name = options["env_var_name"]
+
+            if options.get("binary", False):
+                _sanitize_and_add_entry_to_result(
+                    env_var_name, base64.b64encode(response.payload.data)
+                )
+            else:
+                _sanitize_and_add_entry_to_result(env_var_name, payload_str)
+
+        return result

metaflow/plugins/gcp/gs_storage_client_factory.py

@@ -8,7 +8,7 @@ def _get_cache_key():
     return os.getpid(), threading.get_ident()
 
 
-def get_gs_storage_client():
+def _get_gs_storage_client_default():
     cache_key = _get_cache_key()
     if cache_key not in _client_cache:
         from google.cloud import storage
@@ -19,3 +19,54 @@ def get_gs_storage_client():
             credentials=credentials, project=project_id
         )
     return _client_cache[cache_key]
+
+
+class GcpDefaultClientProvider(object):
+    name = "gcp-default"
+
+    @staticmethod
+    def get_gs_storage_client(*args, **kwargs):
+        return _get_gs_storage_client_default()
+
+    @staticmethod
+    def get_credentials(scopes, *args, **kwargs):
+        import google.auth
+
+        return google.auth.default(scopes=scopes)
+
+
+cached_provider_class = None
+
+
+def get_gs_storage_client():
+    global cached_provider_class
+    if cached_provider_class is None:
+        from metaflow.metaflow_config import DEFAULT_GCP_CLIENT_PROVIDER
+        from metaflow.plugins import GCP_CLIENT_PROVIDERS
+
+        for p in GCP_CLIENT_PROVIDERS:
+            if p.name == DEFAULT_GCP_CLIENT_PROVIDER:
+                cached_provider_class = p
+                break
+        else:
+            raise ValueError(
+                "Cannot find GCP Client provider %s" % DEFAULT_GCP_CLIENT_PROVIDER
+            )
+    return cached_provider_class.get_gs_storage_client()
+
+
+def get_credentials(scopes, *args, **kwargs):
+    global cached_provider_class
+    if cached_provider_class is None:
+        from metaflow.metaflow_config import DEFAULT_GCP_CLIENT_PROVIDER
+        from metaflow.plugins import GCP_CLIENT_PROVIDERS
+
+        for p in GCP_CLIENT_PROVIDERS:
+            if p.name == DEFAULT_GCP_CLIENT_PROVIDER:
+                cached_provider_class = p
+                break
+        else:
+            raise ValueError(
+                "Cannot find GCP Client provider %s" % DEFAULT_GCP_CLIENT_PROVIDER
+            )
+    return cached_provider_class.get_credentials(scopes, *args, **kwargs)

metaflow/plugins/kubernetes/kubernetes.py

@@ -3,9 +3,9 @@ import math
 import os
 import re
 import shlex
+import copy
 import time
 from typing import Dict, List, Optional
-import uuid
 from uuid import uuid4
 
 from metaflow import current, util
@@ -27,8 +27,11 @@ from metaflow.metaflow_config import (
     DATASTORE_SYSROOT_S3,
     DATATOOLS_S3ROOT,
     DEFAULT_AWS_CLIENT_PROVIDER,
+    DEFAULT_GCP_CLIENT_PROVIDER,
     DEFAULT_METADATA,
     DEFAULT_SECRETS_BACKEND_TYPE,
+    GCP_SECRET_MANAGER_PREFIX,
+    AZURE_KEY_VAULT_PREFIX,
     KUBERNETES_FETCH_EC2_METADATA,
     KUBERNETES_LABELS,
     KUBERNETES_SANDBOX_INIT_SCRIPT,
@@ -66,6 +69,12 @@ class KubernetesKilledException(MetaflowException):
     headline = "Kubernetes Batch job killed"
 
 
+def _extract_labels_and_annotations_from_job_spec(job_spec):
+    annotations = job_spec.template.metadata.annotations
+    labels = job_spec.template.metadata.labels
+    return copy.copy(annotations), copy.copy(labels)
+
+
 class Kubernetes(object):
     def __init__(
         self,
@@ -140,9 +149,64 @@ class Kubernetes(object):
         return shlex.split('bash -c "%s"' % cmd_str)
 
     def launch_job(self, **kwargs):
-        self._job = self.create_job(**kwargs).execute()
+        if (
+            "num_parallel" in kwargs
+            and kwargs["num_parallel"]
+            and int(kwargs["num_parallel"]) > 0
+        ):
+            job = self.create_job_object(**kwargs)
+            spec = job.create_job_spec()
+            # `kwargs["step_cli"]` is setting `ubf_context` as control to ALL pods.
+            # This will be modified by the KubernetesJobSet object
+            annotations, labels = _extract_labels_and_annotations_from_job_spec(spec)
+            self._job = self.create_jobset(
+                job_spec=spec,
+                run_id=kwargs["run_id"],
+                step_name=kwargs["step_name"],
+                task_id=kwargs["task_id"],
+                namespace=kwargs["namespace"],
+                env=kwargs["env"],
+                num_parallel=kwargs["num_parallel"],
+                port=kwargs["port"],
+                annotations=annotations,
+                labels=labels,
+            ).execute()
+        else:
+            kwargs["name_pattern"] = "t-{uid}-".format(uid=str(uuid4())[:8])
+            self._job = self.create_job_object(**kwargs).create().execute()
+
+    def create_jobset(
+        self,
+        job_spec=None,
+        run_id=None,
+        step_name=None,
+        task_id=None,
+        namespace=None,
+        env=None,
+        num_parallel=None,
+        port=None,
+        annotations=None,
+        labels=None,
+    ):
+        if env is None:
+            env = {}
+
+        _prefix = str(uuid4())[:6]
+        js = KubernetesClient().jobset(
+            name="js-%s" % _prefix,
+            run_id=run_id,
+            task_id=task_id,
+            step_name=step_name,
+            namespace=namespace,
+            labels=self._get_labels(labels),
+            annotations=annotations,
+            num_parallel=num_parallel,
+            job_spec=job_spec,
+            port=port,
+        )
+        return js
 
-    def create_job(
+    def create_job_object(
         self,
         flow_name,
         run_id,
@@ -176,14 +240,15 @@ class Kubernetes(object):
         labels=None,
         shared_memory=None,
         port=None,
+        name_pattern=None,
+        num_parallel=None,
     ):
         if env is None:
             env = {}
-
         job = (
             KubernetesClient()
             .job(
-                generate_name="t-{uid}-".format(uid=str(uuid4())[:8]),
+                generate_name=name_pattern,
                 namespace=namespace,
                 service_account=service_account,
                 secrets=secrets,
@@ -217,6 +282,7 @@ class Kubernetes(object):
                 persistent_volume_claims=persistent_volume_claims,
                 shared_memory=shared_memory,
                 port=port,
+                num_parallel=num_parallel,
             )
             .environment_variable("METAFLOW_CODE_SHA", code_package_sha)
             .environment_variable("METAFLOW_CODE_URL", code_package_url)
@@ -243,10 +309,19 @@ class Kubernetes(object):
             .environment_variable(
                 "METAFLOW_DEFAULT_AWS_CLIENT_PROVIDER", DEFAULT_AWS_CLIENT_PROVIDER
             )
+            .environment_variable(
+                "METAFLOW_DEFAULT_GCP_CLIENT_PROVIDER", DEFAULT_GCP_CLIENT_PROVIDER
+            )
             .environment_variable(
                 "METAFLOW_AWS_SECRETS_MANAGER_DEFAULT_REGION",
                 AWS_SECRETS_MANAGER_DEFAULT_REGION,
             )
+            .environment_variable(
+                "METAFLOW_GCP_SECRET_MANAGER_PREFIX", GCP_SECRET_MANAGER_PREFIX
+            )
+            .environment_variable(
+                "METAFLOW_AZURE_KEY_VAULT_PREFIX", AZURE_KEY_VAULT_PREFIX
+            )
             .environment_variable("METAFLOW_S3_ENDPOINT_URL", S3_ENDPOINT_URL)
             .environment_variable(
                 "METAFLOW_AZURE_STORAGE_BLOB_SERVICE_ENDPOINT",
@@ -332,6 +407,9 @@ class Kubernetes(object):
             .label("app.kubernetes.io/part-of", "metaflow")
         )
 
+        return job
+
+    def create_k8sjob(self, job):
         return job.create()
 
     def wait(self, stdout_location, stderr_location, echo=None):
@@ -366,7 +444,7 @@ class Kubernetes(object):
                     t = time.time()
                 time.sleep(update_delay(time.time() - start_time))
 
-        prefix = b"[%s] " % util.to_bytes(self._job.id)
+        _make_prefix = lambda: b"[%s] " % util.to_bytes(self._job.id)
 
         stdout_tail = get_log_tailer(stdout_location, self._datastore.TYPE)
         stderr_tail = get_log_tailer(stderr_location, self._datastore.TYPE)
@@ -376,7 +454,7 @@ class Kubernetes(object):
 
         # 2) Tail logs until the job has finished
         tail_logs(
-            prefix=prefix,
+            prefix=_make_prefix(),
             stdout_tail=stdout_tail,
             stderr_tail=stderr_tail,
             echo=echo,
@@ -392,7 +470,6 @@ class Kubernetes(object):
         #        exists prior to calling S3Tail and note the user about
         #        truncated logs if it doesn't.
         # TODO : For hard crashes, we can fetch logs from the pod.
-
         if self._job.has_failed:
             exit_code, reason = self._job.reason
             msg = next(

metaflow/plugins/kubernetes/kubernetes_cli.py

@@ -7,11 +7,17 @@ from metaflow import JSONTypeClass, util
 from metaflow._vendor import click
 from metaflow.exception import METAFLOW_EXIT_DISALLOW_RETRY, CommandException
 from metaflow.metadata.util import sync_local_metadata_from_datastore
+from metaflow.unbounded_foreach import UBF_CONTROL, UBF_TASK
 from metaflow.metaflow_config import DATASTORE_LOCAL_DIR, KUBERNETES_LABELS
 from metaflow.mflog import TASK_LOG_SOURCE
 import metaflow.tracing as tracing
 
-from .kubernetes import Kubernetes, KubernetesKilledException, parse_kube_keyvalue_list
+from .kubernetes import (
+    Kubernetes,
+    KubernetesKilledException,
+    parse_kube_keyvalue_list,
+    KubernetesException,
+)
 from .kubernetes_decorator import KubernetesDecorator
 
 
@@ -109,6 +115,15 @@ def kubernetes():
 )
 @click.option("--shared-memory", default=None, help="Size of shared memory in MiB")
 @click.option("--port", default=None, help="Port number to expose from the container")
+@click.option(
+    "--ubf-context", default=None, type=click.Choice([None, UBF_CONTROL, UBF_TASK])
+)
+@click.option(
+    "--num-parallel",
+    default=None,
+    type=int,
+    help="Number of parallel nodes to run as a multi-node job.",
+)
 @click.pass_context
 def step(
     ctx,
@@ -136,6 +151,7 @@ def step(
     tolerations=None,
     shared_memory=None,
     port=None,
+    num_parallel=None,
     **kwargs
 ):
     def echo(msg, stream="stderr", job_id=None, **kwargs):
@@ -167,6 +183,12 @@ def step(
         kwargs["input_paths"] = "".join("${%s}" % s for s in split_vars.keys())
         env.update(split_vars)
 
+    if num_parallel is not None and num_parallel <= 1:
+        raise KubernetesException(
+            "Using @parallel with `num_parallel` <= 1 is not supported with Kubernetes. "
+            "Please set the value of `num_parallel` to be greater than 1."
+        )
+
     # Set retry policy.
     retry_count = int(kwargs.get("retry_count", 0))
     retry_deco = [deco for deco in node.decorators if deco.name == "retry"]
@@ -251,6 +273,7 @@ def step(
                 tolerations=tolerations,
                 shared_memory=shared_memory,
                 port=port,
+                num_parallel=num_parallel,
             )
     except Exception as e:
         traceback.print_exc(chain=False)

metaflow/plugins/kubernetes/kubernetes_client.py

@@ -4,7 +4,7 @@ import time
 
 from metaflow.exception import MetaflowException
 
-from .kubernetes_job import KubernetesJob
+from .kubernetes_job import KubernetesJob, KubernetesJobSet
 
 
 CLIENT_REFRESH_INTERVAL_SECONDS = 300
@@ -61,5 +61,8 @@ class KubernetesClient(object):
 
         return self._client
 
+    def jobset(self, **kwargs):
+        return KubernetesJobSet(self, **kwargs)
+
     def job(self, **kwargs):
         return KubernetesJob(self, **kwargs)

metaflow/plugins/kubernetes/kubernetes_decorator.py

@@ -32,6 +32,8 @@ from metaflow.sidecar import Sidecar
 
 from ..aws.aws_utils import get_docker_registry, get_ec2_instance_metadata
 from .kubernetes import KubernetesException, parse_kube_keyvalue_list
+from metaflow.unbounded_foreach import UBF_CONTROL
+from .kubernetes_jobsets import TaskIdConstructor
 
 try:
     unicode
@@ -239,11 +241,15 @@ class KubernetesDecorator(StepDecorator):
                 "Kubernetes. Please use one or the other.".format(step=step)
             )
 
-        for deco in decos:
-            if getattr(deco, "IS_PARALLEL", False):
-                raise KubernetesException(
-                    "@kubernetes does not support parallel execution currently."
+        if any([deco.name == "parallel" for deco in decos]) and any(
+            [deco.name == "catch" for deco in decos]
+        ):
+            raise MetaflowException(
+                "Step *{step}* contains a @parallel decorator "
+                "with the @catch decorator. @catch is not supported with @parallel on Kubernetes.".format(
+                    step=step
                 )
+            )
 
         # Set run time limit for the Kubernetes job.
         self.run_time_limit = get_run_time_limit_for_task(decos)
@@ -421,6 +427,10 @@ class KubernetesDecorator(StepDecorator):
                 "METAFLOW_KUBERNETES_SERVICE_ACCOUNT_NAME"
             ]
             meta["kubernetes-node-ip"] = os.environ["METAFLOW_KUBERNETES_NODE_IP"]
+            if os.environ.get("METAFLOW_KUBERNETES_JOBSET_NAME"):
+                meta["kubernetes-jobset-name"] = os.environ[
+                    "METAFLOW_KUBERNETES_JOBSET_NAME"
+                ]
 
             # TODO (savin): Introduce equivalent support for Microsoft Azure and
             #               Google Cloud Platform
@@ -453,6 +463,24 @@ class KubernetesDecorator(StepDecorator):
             self._save_logs_sidecar = Sidecar("save_logs_periodically")
             self._save_logs_sidecar.start()
 
+        num_parallel = None
+        if hasattr(flow, "_parallel_ubf_iter"):
+            num_parallel = flow._parallel_ubf_iter.num_parallel
+
+        if num_parallel and num_parallel >= 1 and ubf_context == UBF_CONTROL:
+            control_task_id, worker_task_ids = TaskIdConstructor.join_step_task_ids(
+                num_parallel
+            )
+            mapper_task_ids = [control_task_id] + worker_task_ids
+            flow._control_mapper_tasks = [
+                "%s/%s/%s" % (run_id, step_name, mapper_task_id)
+                for mapper_task_id in mapper_task_ids
+            ]
+            flow._control_task_is_mapper_zero = True
+
+        if num_parallel and num_parallel > 1:
+            _setup_multinode_environment()
+
     def task_finished(
         self, step_name, flow, graph, is_task_ok, retry_count, max_retries
     ):
@@ -486,3 +514,20 @@ class KubernetesDecorator(StepDecorator):
             cls.package_url, cls.package_sha = flow_datastore.save_data(
                 [package.blob], len_hint=1
             )[0]
+
+
+def _setup_multinode_environment():
+    import socket
+
+    os.environ["MF_PARALLEL_MAIN_IP"] = socket.gethostbyname(os.environ["MASTER_ADDR"])
+    os.environ["MF_PARALLEL_NUM_NODES"] = os.environ["WORLD_SIZE"]
+    if os.environ.get("CONTROL_INDEX") is not None:
+        os.environ["MF_PARALLEL_NODE_INDEX"] = str(0)
+    elif os.environ.get("WORKER_REPLICA_INDEX") is not None:
+        os.environ["MF_PARALLEL_NODE_INDEX"] = str(
+            int(os.environ["WORKER_REPLICA_INDEX"]) + 1
+        )
+    else:
+        raise MetaflowException(
+            "Jobset related ENV vars called $CONTROL_INDEX or $WORKER_REPLICA_INDEX not found"
+        )