ob-metaflow 2.11.15.3__py2.py3-none-any.whl → 2.11.16.1__py2.py3-none-any.whl

metaflow/__init__.py

@@ -143,6 +143,9 @@ from .client import (
     DataArtifact,
 )
 
+# Import data class within tuple_util but not introduce new symbols.
+from . import tuple_util
+
 __version_addl__ = []
 _ext_debug("Loading top-level modules")
 for m in _tl_modules:

metaflow/clone_util.py

@@ -66,6 +66,12 @@ def clone_task_helper(
                 type="attempt",
                 tags=metadata_tags,
             ),
+            MetaDatum(
+                field="attempt_ok",
+                value="True",  # During clone, the task is always considered successful.
+                type="internal_attempt_status",
+                tags=metadata_tags,
+            ),
         ],
     )
     output.done()

metaflow/extension_support/plugins.py

@@ -179,8 +179,8 @@ _plugin_categories = {
     "metadata_provider": lambda x: x.TYPE,
     "datastore": lambda x: x.TYPE,
     "secrets_provider": lambda x: x.TYPE,
-    "azure_client_provider": lambda x: x.name,
     "gcp_client_provider": lambda x: x.name,
+    "azure_client_provider": lambda x: x.name,
     "sidecar": None,
     "logging_sidecar": None,
     "monitor_sidecar": None,

metaflow/metaflow_config.py

@@ -227,6 +227,8 @@ DEFAULT_CONTAINER_REGISTRY = from_conf("DEFAULT_CONTAINER_REGISTRY")
 INCLUDE_FOREACH_STACK = from_conf("INCLUDE_FOREACH_STACK", False)
 # Maximum length of the foreach value string to be stored in each ForeachFrame.
 MAXIMUM_FOREACH_VALUE_CHARS = from_conf("MAXIMUM_FOREACH_VALUE_CHARS", 30)
+# The default runtime limit (In seconds) of jobs launched by any compute provider. Default of 5 days.
+DEFAULT_RUNTIME_LIMIT = from_conf("DEFAULT_RUNTIME_LIMIT", 5 * 24 * 60 * 60)
 
 ###
 # Organization customizations
@@ -327,8 +329,6 @@ KUBERNETES_CONTAINER_REGISTRY = from_conf(
 )
 # Toggle for trying to fetch EC2 instance metadata
 KUBERNETES_FETCH_EC2_METADATA = from_conf("KUBERNETES_FETCH_EC2_METADATA", False)
-# Default port number to open on the pods
-KUBERNETES_PORT = from_conf("KUBERNETES_PORT", None)
 # Shared memory in MB to use for this step
 KUBERNETES_SHARED_MEMORY = from_conf("KUBERNETES_SHARED_MEMORY", None)
 # Default port number to open on the pods
@@ -338,10 +338,12 @@ KUBERNETES_CPU = from_conf("KUBERNETES_CPU", None)
 KUBERNETES_MEMORY = from_conf("KUBERNETES_MEMORY", None)
 KUBERNETES_DISK = from_conf("KUBERNETES_DISK", None)
 
-
 ARGO_WORKFLOWS_KUBERNETES_SECRETS = from_conf("ARGO_WORKFLOWS_KUBERNETES_SECRETS", "")
 ARGO_WORKFLOWS_ENV_VARS_TO_SKIP = from_conf("ARGO_WORKFLOWS_ENV_VARS_TO_SKIP", "")
 
+KUBERNETES_JOBSET_GROUP = from_conf("KUBERNETES_JOBSET_GROUP", "jobset.x-k8s.io")
+KUBERNETES_JOBSET_VERSION = from_conf("KUBERNETES_JOBSET_VERSION", "v1alpha2")
+
 ##
 # Argo Events Configuration
 ##

metaflow/metaflow_environment.py

@@ -91,7 +91,7 @@ class MetaflowEnvironment(object):
         if datastore_type == "s3":
             return (
                 '%s -m awscli ${METAFLOW_S3_ENDPOINT_URL:+--endpoint-url=\\"${METAFLOW_S3_ENDPOINT_URL}\\"} '
-                + "s3 cp %s job.tar"
+                + "s3 cp %s job.tar >/dev/null"
             ) % (self._python(), code_package_url)
         elif datastore_type == "azure":
             from .plugins.azure.azure_utils import parse_azure_full_path
@@ -119,9 +119,9 @@ class MetaflowEnvironment(object):
             )
 
     def _get_install_dependencies_cmd(self, datastore_type):
-        cmds = ["%s -m pip install requests" % self._python()]
+        cmds = ["%s -m pip install requests -qqq" % self._python()]
         if datastore_type == "s3":
-            cmds.append("%s -m pip install awscli boto3" % self._python())
+            cmds.append("%s -m pip install awscli boto3 -qqq" % self._python())
         elif datastore_type == "azure":
             cmds.append(
                 "%s -m pip install azure-identity azure-storage-blob azure-keyvault-secrets simple-azure-blob-downloader -qqq"

metaflow/plugins/__init__.py

@@ -131,14 +131,14 @@ SECRETS_PROVIDERS_DESC = [
     ),
 ]
 
-AZURE_CLIENT_PROVIDERS_DESC = [
-    ("azure-default", ".azure.azure_credential.AzureDefaultClientProvider")
-]
-
 GCP_CLIENT_PROVIDERS_DESC = [
     ("gcp-default", ".gcp.gs_storage_client_factory.GcpDefaultClientProvider")
 ]
 
+AZURE_CLIENT_PROVIDERS_DESC = [
+    ("azure-default", ".azure.azure_credential.AzureDefaultClientProvider")
+]
+
 
 process_plugins(globals())
 

metaflow/plugins/azure/azure_secret_manager_secrets_provider.py

@@ -80,7 +80,7 @@ class AzureKeyVaultSecretsProvider(SecretsProvider):
         try:
             parsed_vault_url = urlparse(secret_id)
         except ValueError:
-            print(f"invalid vault url", file=sys.stderr)
+            print("invalid vault url", file=sys.stderr)
             return False
         hostname = parsed_vault_url.netloc
 
@@ -94,34 +94,34 @@ class AzureKeyVaultSecretsProvider(SecretsProvider):
         if not k_v_domain_found:
             # the secret_id started with https:// however the key_vault_domains
             # were not present in the secret_id which means
-            raise MetaflowAzureKeyVaultBadVault(f"bad key vault domain {secret_id}")
+            raise MetaflowAzureKeyVaultBadVault("bad key vault domain %s" % secret_id)
 
         # given the secret_id seems to have a valid key vault domain
         # lets verify that the vault name corresponds to its regex.
         vault_name = hostname[: -len(actual_k_v_domain)]
         # verify the vault name pattern
         if not self._is_valid_vault_name(vault_name):
-            raise MetaflowAzureKeyVaultBadVault(f"bad key vault name {vault_name}")
+            raise MetaflowAzureKeyVaultBadVault("bad key vault name %s" % vault_name)
 
         path_parts = parsed_vault_url.path.strip("/").split("/")
         total_path_parts = len(path_parts)
         if total_path_parts < 2 or total_path_parts > 3:
             raise MetaflowAzureKeyVaultBadSecretPath(
-                f"bad secret uri path {path_parts}"
+                "bad secret uri path %s" % path_parts
             )
 
         object_type = path_parts[0]
         if not self._is_valid_object_type(object_type):
-            raise MetaflowAzureKeyVaultBadSecretType(f"bad secret type {object_type}")
+            raise MetaflowAzureKeyVaultBadSecretType("bad secret type %s" % object_type)
 
         secret_name = path_parts[1]
         if not self._is_valid_secret_name(secret_name=secret_name):
-            raise MetaflowAzureKeyVaultBadSecretName(f"bad secret name {secret_name}")
+            raise MetaflowAzureKeyVaultBadSecretName("bad secret name %s" % secret_name)
 
         if total_path_parts == 3:
             if not self._is_valid_object_version(path_parts[2]):
                 raise MetaflowAzureKeyVaultBadSecretVersion(
-                    f"bad secret version {path_parts[2]}"
+                    "bad secret version %s" % path_parts[2]
                 )
 
         return True
@@ -139,10 +139,11 @@ class AzureKeyVaultSecretsProvider(SecretsProvider):
         # must be set.
         if not AZURE_KEY_VAULT_PREFIX:
             raise ValueError(
-                f"cannot use simple secret id without setting METAFLOW_AZURE_KEY_VAULT_PREFIX. {AZURE_KEY_VAULT_PREFIX}"
+                "cannot use simple secret id without setting METAFLOW_AZURE_KEY_VAULT_PREFIX. %s"
+                % AZURE_KEY_VAULT_PREFIX
             )
         domain = AZURE_KEY_VAULT_PREFIX.rstrip("/")
-        full_secret = f"{domain}/secrets/{secret_id}"
+        full_secret = "%s/secrets/%s" % (domain, secret_id)
         if not self._is_secret_id_fully_qualified_url(full_secret):
             return False
 
@@ -186,29 +187,32 @@ class AzureKeyVaultSecretsProvider(SecretsProvider):
 
         # if the secret_id is None/empty/does not start with https then return false
         if secret_id is None or secret_id == "":
-            raise MetaflowAzureKeyVaultBadSecret(f"empty secret id is not supported")
+            raise MetaflowAzureKeyVaultBadSecret("empty secret id is not supported")
 
         # check if the passed in secret is a short-form ( #3/#4 in the above comment)
         if not secret_id.startswith("https://"):
             # check if the secret_id is of form `secret_name` OR `secret_name/secret_version`
             if not self._is_partial_secret_valid(secret_id=secret_id):
                 raise MetaflowAzureKeyVaultBadSecret(
-                    f"unsupported partial secret {secret_id}"
+                    "unsupported partial secret %s" % secret_id
                 )
 
             domain = AZURE_KEY_VAULT_PREFIX.rstrip("/")
-            full_secret = f"{domain}/secrets/{secret_id}"
+            full_secret = "%s/secrets/%s" % (domain, secret_id)
 
         # if the secret id is passed as a URL - then check if the url is fully qualified
         if secret_id.startswith("https://"):
             if not self._is_secret_id_fully_qualified_url(secret_id=secret_id):
-                raise MetaflowException(f"unsupported secret {secret_id}")
+                raise MetaflowException("unsupported secret %s" % secret_id)
             full_secret = secret_id
 
         # at this point I know that the secret URL is good so we can start creating the Secret Client
         az_credentials = create_cacheable_azure_credential()
         res = urlparse(full_secret)
-        az_vault_url = f"{res.scheme}://{res.netloc}"  # https://myvault.vault.azure.net
+        az_vault_url = "%s://%s" % (
+            res.scheme,
+            res.netloc,
+        )  # https://myvault.vault.azure.net
         secret_data = res.path.strip("/").split("/")[1:]
         secret_name = secret_data[0]
         secret_version = None

metaflow/plugins/datatools/s3/s3.py

@@ -21,7 +21,6 @@ from metaflow.metaflow_config import (
     TEMPDIR,
 )
 from metaflow.util import (
-    namedtuple_with_defaults,
     is_stringish,
     to_bytes,
     to_unicode,
@@ -29,6 +28,7 @@ from metaflow.util import (
     url_quote,
     url_unquote,
 )
+from metaflow.tuple_util import namedtuple_with_defaults
 from metaflow.exception import MetaflowException
 from metaflow.debug import debug
 import metaflow.tracing as tracing

metaflow/plugins/gcp/__init__.py

@@ -1 +1 @@
-from .gs_storage_client_factory import get_credentials
+from .gs_storage_client_factory import get_credentials

metaflow/plugins/gcp/gcp_secret_manager_secrets_provider.py

@@ -89,7 +89,9 @@ class GcpSecretManagerSecretsProvider(SecretsProvider):
 
         # 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_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
@@ -105,20 +107,23 @@ class GcpSecretManagerSecretsProvider(SecretsProvider):
             env_var_name = secret_id
             if not GCP_SECRET_MANAGER_PREFIX:
                 raise ValueError(
-                    f"Cannot use simple secret_id without setting METAFLOW_GCP_SECRET_MANAGER_PREFIX. {GCP_SECRET_MANAGER_PREFIX}"
+                    "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 = f"{GCP_SECRET_MANAGER_PREFIX}{secret_id}"
+                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 = (
-                    f"{GCP_SECRET_MANAGER_PREFIX}{secret_id}/versions/latest"
+                full_secret_name = "%s%s/versions/latest" % (
+                    GCP_SECRET_MANAGER_PREFIX,
+                    secret_id,
                 )
         else:
             raise ValueError(
-                f"Invalid secret_id: {secret_id}. Must be either a full path or a simple string."
+                "Invalid secret_id: %s. Must be either a full path or a simple string."
+                % secret_id
             )
 
         result = {}

metaflow/plugins/kubernetes/kubernetes.py

@@ -3,10 +3,9 @@ import math
 import os
 import re
 import shlex
-import time
 import copy
+import time
 from typing import Dict, List, Optional
-import uuid
 from uuid import uuid4
 
 from metaflow import current, util
@@ -70,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,
@@ -144,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,
@@ -178,19 +238,17 @@ class Kubernetes(object):
         persistent_volume_claims=None,
         tolerations=None,
         labels=None,
-        annotations=None,
-        num_parallel=0,
-        attrs={},
         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,
@@ -222,10 +280,9 @@ class Kubernetes(object):
                 tmpfs_size=tmpfs_size,
                 tmpfs_path=tmpfs_path,
                 persistent_volume_claims=persistent_volume_claims,
-                num_parallel=num_parallel,
-                attrs=attrs,
                 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)
@@ -288,7 +345,6 @@ class Kubernetes(object):
             # see get_datastore_root_from_config in datastore/local.py).
         )
 
-        self.num_parallel = num_parallel
         # Temporary passing of *some* environment variables. Do not rely on this
         # mechanism as it will be removed in the near future
         for k, v in config_values():
@@ -351,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):
@@ -364,7 +423,7 @@ class Kubernetes(object):
             sigmoid = 1.0 / (1.0 + math.exp(-0.01 * secs_since_start + 9.0))
             return 0.5 + sigmoid * 30.0
 
-        def wait_for_launch(job, child_jobs):
+        def wait_for_launch(job):
             status = job.status
             echo(
                 "Task is starting (%s)..." % status,
@@ -374,56 +433,28 @@ class Kubernetes(object):
             t = time.time()
             start_time = time.time()
             while job.is_waiting:
-                # new_status = job.status
-                if status != job.status or (time.time() - t) > 30:
-                    if not child_jobs:
-                        child_statuses = ""
-                    else:
-                        status_keys = set(
-                            [child_job.status for child_job in child_jobs]
-                        )
-                        status_counts = [
-                            (
-                                status,
-                                len(
-                                    [
-                                        child_job.status == status
-                                        for child_job in child_jobs
-                                    ]
-                                ),
-                            )
-                            for status in status_keys
-                        ]
-                        child_statuses = " (parallel node status: [{}])".format(
-                            ", ".join(
-                                [
-                                    "{}:{}".format(status, num)
-                                    for (status, num) in sorted(status_counts)
-                                ]
-                            )
-                        )
-
-                    status = job.status
+                new_status = job.status
+                if status != new_status or (time.time() - t) > 30:
+                    status = new_status
                     echo(
-                        "Task is starting (status %s)... %s" % (status, child_statuses),
+                        "Task is starting (%s)..." % status,
                         "stderr",
                         job_id=job.id,
                     )
                     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)
 
-        child_jobs = []
         # 1) Loop until the job has started
-        wait_for_launch(self._job, child_jobs)
+        wait_for_launch(self._job)
 
         # 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,
@@ -439,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
 
 
@@ -107,27 +113,17 @@ def kubernetes():
     type=JSONTypeClass(),
     multiple=False,
 )
+@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(
-    "--labels",
-    default=None,
-    type=JSONTypeClass(),
-    multiple=False,
-)
-@click.option(
-    "--annotations",
-    default=None,
-    type=JSONTypeClass(),
-    multiple=False,
+    "--ubf-context", default=None, type=click.Choice([None, UBF_CONTROL, UBF_TASK])
 )
-@click.option("--ubf-context", default=None, type=click.Choice([None, "ubf_control"]))
 @click.option(
     "--num-parallel",
-    default=0,
+    default=None,
     type=int,
     help="Number of parallel nodes to run as a multi-node job.",
 )
-@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.pass_context
 def step(
     ctx,
@@ -153,11 +149,9 @@ def step(
     run_time_limit=None,
     persistent_volume_claims=None,
     tolerations=None,
-    labels=None,
-    annotations=None,
-    num_parallel=None,
     shared_memory=None,
     port=None,
+    num_parallel=None,
     **kwargs
 ):
     def echo(msg, stream="stderr", job_id=None, **kwargs):
@@ -189,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"]
@@ -203,17 +203,11 @@ def step(
         )
         time.sleep(minutes_between_retries * 60)
 
-    step_args = " ".join(util.dict_to_cli_options(kwargs))
-    num_parallel = num_parallel or 0
-    if num_parallel and num_parallel > 1:
-        # For multinode, we need to add a placeholder that can be mutated by the caller
-        step_args += " [multinode-args]"
-
     step_cli = "{entrypoint} {top_args} step {step} {step_args}".format(
         entrypoint="%s -u %s" % (executable, os.path.basename(sys.argv[0])),
         top_args=" ".join(util.dict_to_cli_options(ctx.parent.parent.params)),
         step=step_name,
-        step_args=step_args,
+        step_args=" ".join(util.dict_to_cli_options(kwargs)),
     )
 
     # Set log tailing.
@@ -239,10 +233,6 @@ def step(
                 ),
             )
 
-    attrs = {
-        "metaflow.task_id": kwargs["task_id"],
-        "requires_passwordless_ssh": any([getattr(deco, "requires_passwordless_ssh", False) for deco in node.decorators]),
-    }
     try:
         kubernetes = Kubernetes(
             datastore=ctx.obj.flow_datastore,
@@ -281,12 +271,9 @@ def step(
                 env=env,
                 persistent_volume_claims=persistent_volume_claims,
                 tolerations=tolerations,
-                labels=labels,
-                annotations=annotations,
-                num_parallel=num_parallel,
                 shared_memory=shared_memory,
                 port=port,
-                attrs=attrs,
+                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)