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

metaflow/plugins/kubernetes/kubernetes_job.py

@@ -2,18 +2,23 @@ import json
 import math
 import random
 import time
-import os
-import socket
 import copy
-
+import sys
+from metaflow.tracing import inject_tracing_vars
 from metaflow.exception import MetaflowException
 from metaflow.metaflow_config import KUBERNETES_SECRETS
+from metaflow.unbounded_foreach import UBF_CONTROL, UBF_TASK
 
 CLIENT_REFRESH_INTERVAL_SECONDS = 300
+from .kubernetes_jobsets import (
+    KubernetesJobSet,  # We need this import for Kubernetes Client.
+)
+
 
 class KubernetesJobException(MetaflowException):
     headline = "Kubernetes job error"
 
+
 # Implements truncated exponential backoff from
 # https://cloud.google.com/storage/docs/retry-strategy#exponential-backoff
 def k8s_retry(deadline_seconds=60, max_backoff=32):
@@ -56,19 +61,7 @@ class KubernetesJob(object):
         self._client = client
         self._kwargs = kwargs
 
-    def create(self):
-        # A discerning eye would notice and question the choice of using the
-        # V1Job construct over the V1Pod construct given that we don't rely much
-        # on any of the V1Job semantics. The major reasons at the moment are -
-        #     1. It makes the Kubernetes UIs (Octant, Lens) a bit easier on
-        #        the eyes, although even that can be questioned.
-        #     2. AWS Step Functions, at the moment (Apr' 22) only supports
-        #        executing Jobs and not Pods as part of it's publicly declared
-        #        API. When we ship the AWS Step Functions integration with EKS,
-        #        it will hopefully lessen our workload.
-        #
-        # Note: This implementation ensures that there is only one unique Pod
-        # (unique UID) per Metaflow task attempt.
+    def create_job_spec(self):
         client = self._client.get()
 
         # tmpfs variables
@@ -80,529 +73,103 @@ class KubernetesJob(object):
             if self._kwargs["shared_memory"]
             else None
         )
-
-        jobset_name = "js-%s" % self._kwargs["attrs"]["metaflow.task_id"].split('-')[-1]
-        main_job_name = "control"
-        main_job_index = 0
-        main_pod_index = 0
-        subdomain = jobset_name
-        master_port = int(self._kwargs['port']) if self._kwargs['port'] else None
-        shared_memory = int(self._kwargs['shared_memory']) if self._kwargs['shared_memory'] else None
-
-        passwordless_ssh = self._kwargs["attrs"]["requires_passwordless_ssh"]
-        if passwordless_ssh:
-            passwordless_ssh_service_name = subdomain
-            passwordless_ssh_service_selector = {
-                "passwordless-ssh-jobset": "true"
-            }
-        else:
-            passwordless_ssh_service_name = None
-            passwordless_ssh_service_selector = {}
-
-        fqdn_suffix = "%s.svc.cluster.local" % self._kwargs["namespace"]
-        jobset_main_addr = "%s-%s-%s-%s.%s.%s" % (
-            jobset_name,
-            main_job_name,
-            main_job_index,
-            main_pod_index,
-            subdomain,
-            fqdn_suffix,
-        )
-
-        def _install_jobset(
-            repo_url="https://github.com/kubernetes-sigs/jobset",
-            python_sdk_path="jobset/sdk/python",
-        ):
-
-            # TODO (Eddie): Remove this and suggest to user.
-
-            import subprocess
-            import tempfile
-            import shutil
-            import os
-
-            with open(os.devnull, "wb") as devnull:
-                cwd = os.getcwd()
-                tmp_dir = tempfile.mkdtemp()
-                os.chdir(tmp_dir)
-                subprocess.check_call(
-                    ["git", "clone", repo_url], stdout=devnull, stderr=subprocess.STDOUT
-                )
-                tmp_python_sdk_path = os.path.join(tmp_dir, python_sdk_path)
-                os.chdir(tmp_python_sdk_path)
-                subprocess.check_call(
-                    ["pip", "install", "."], stdout=devnull, stderr=subprocess.STDOUT
-                )
-                os.chdir(cwd)
-                shutil.rmtree(tmp_dir)
-
-        def _get_passwordless_ssh_service():
-
-            return client.V1Service(
-                api_version="v1",
-                kind="Service",
-                metadata=client.V1ObjectMeta(
-                    name=passwordless_ssh_service_name,
-                    namespace=self._kwargs["namespace"]
-                ),
-                spec=client.V1ServiceSpec(
-                    cluster_ip="None",
-                    internal_traffic_policy="Cluster",
-                    ip_families=["IPv4"],
-                    ip_family_policy="SingleStack",
-                    selector=passwordless_ssh_service_selector,
-                    session_affinity="None",
-                    type="ClusterIP",
-                    ports=[
-                        client.V1ServicePort(
-                            name="control",
-                            port=22,
-                            protocol="TCP",
-                            target_port=22
-                        )
-                    ]
-                )
-            )
-
-        def _get_replicated_job(job_name, parallelism, command):
-            return jobset.models.jobset_v1alpha2_replicated_job.JobsetV1alpha2ReplicatedJob(
-                name=job_name,
-                template=client.V1JobTemplateSpec(
-                    metadata=client.V1ObjectMeta(
-                        annotations=self._kwargs.get("annotations", {}),
-                        labels=self._kwargs.get("labels", {}),
-                        namespace=self._kwargs["namespace"],
-                    ),
-                    spec=client.V1JobSpec(
-                        parallelism=parallelism,  # how many jobs can run at once
-                        completions=parallelism,  # how many Pods the JobSet creates in total
-                        backoff_limit=0,
-                        ttl_seconds_after_finished=7
-                        * 60
-                        * 60
-                        * 24,
-                        template=client.V1PodTemplateSpec(
-                            metadata=client.V1ObjectMeta(
-                                annotations=self._kwargs.get("annotations", {}),
-                                labels={
-                                    **self._kwargs.get("labels", {}),
-                                    **passwordless_ssh_service_selector, # TODO: necessary?
-                                    # TODO: cluster-name, app.kubernetes.io/name necessary?
-                                },
-                                namespace=self._kwargs["namespace"],
-                            ),
-                            spec=client.V1PodSpec(
-                                active_deadline_seconds=self._kwargs[
-                                    "timeout_in_seconds"
-                                ],
-                                containers=[
-                                    client.V1Container(
-                                        command=command,
-                                        ports=[client.V1ContainerPort(container_port=master_port)] if master_port and job_name=="control" else [],
-                                        env=[
-                                            client.V1EnvVar(name=k, value=str(v))
-                                            for k, v in self._kwargs.get(
-                                                "environment_variables", {}
-                                            ).items()
-                                        ]
-                                        + [
-                                            client.V1EnvVar(
-                                                name=k,
-                                                value_from=client.V1EnvVarSource(
-                                                    field_ref=client.V1ObjectFieldSelector(
-                                                        field_path=str(v)
-                                                    )
-                                                ),
-                                            )
-                                            for k, v in {
-                                                "METAFLOW_KUBERNETES_POD_NAMESPACE": "metadata.namespace",
-                                                "METAFLOW_KUBERNETES_POD_NAME": "metadata.name",
-                                                "METAFLOW_KUBERNETES_POD_ID": "metadata.uid",
-                                                "METAFLOW_KUBERNETES_SERVICE_ACCOUNT_NAME": "spec.serviceAccountName",
-                                                "METAFLOW_KUBERNETES_NODE_IP": "status.hostIP",
-                                            }.items()
-                                        ]
-                                        # Mimicking the AWS Batch Multinode env vars.
-                                        + [
-                                            client.V1EnvVar(
-                                                name="MASTER_ADDR",
-                                                value=jobset_main_addr,
-                                            ),
-                                            client.V1EnvVar(
-                                                name="MASTER_PORT",
-                                                value=str(master_port),
-                                            ),
-                                            client.V1EnvVar(
-                                                name="RANK",
-                                                value_from=client.V1EnvVarSource(
-                                                    field_ref=client.V1ObjectFieldSelector(
-                                                        field_path="metadata.annotations['batch.kubernetes.io/job-completion-index']"
-                                                    )
-                                                ),
-                                            ),
-                                            client.V1EnvVar(
-                                                name="WORLD_SIZE",
-                                                value=str(self._kwargs["num_parallel"]),
-                                            ),
-                                            client.V1EnvVar(
-                                                name="PYTHONUNBUFFERED",
-                                                value="0",
-                                            ),
-                                        ],
-                                        env_from=[
-                                            client.V1EnvFromSource(
-                                                secret_ref=client.V1SecretEnvSource(
-                                                    name=str(k),
-                                                    # optional=True
-                                                )
-                                            )
-                                            for k in list(
-                                                self._kwargs.get("secrets", [])
-                                            )
-                                            + KUBERNETES_SECRETS.split(",")
-                                            if k
-                                        ],
-                                        image=self._kwargs["image"],
-                                        image_pull_policy=self._kwargs[
-                                            "image_pull_policy"
-                                        ],
-                                        name=self._kwargs["step_name"].replace(
-                                            "_", "-"
-                                        ),
-                                        resources=client.V1ResourceRequirements(
-                                            requests={
-                                                "cpu": str(self._kwargs["cpu"]),
-                                                "memory": "%sM"
-                                                % str(self._kwargs["memory"]),
-                                                "ephemeral-storage": "%sM"
-                                                % str(self._kwargs["disk"]),
-                                            },
-                                            limits={
-                                                "%s.com/gpu".lower()
-                                                % self._kwargs["gpu_vendor"]: str(
-                                                    self._kwargs["gpu"]
-                                                )
-                                                for k in [0]
-                                                # Don't set GPU limits if gpu isn't specified.
-                                                if self._kwargs["gpu"] is not None
-                                            },
-                                        ),
-                                        volume_mounts=(
-                                            [
-                                                client.V1VolumeMount(
-                                                    mount_path=self._kwargs.get(
-                                                        "tmpfs_path"
-                                                    ),
-                                                    name="tmpfs-ephemeral-volume",
-                                                )
-                                            ]
-                                            if tmpfs_enabled
-                                            else []
-                                        )
-                                        + (
-                                            [
-                                                client.V1VolumeMount(
-                                                    mount_path="/dev/shm",
-                                                    name="dhsm"
-                                                )
-                                            ]
-                                            if shared_memory else []
-                                        )
-                                        + (
-                                            [
-                                                client.V1VolumeMount(
-                                                    mount_path=path, name=claim
-                                                )
-                                                for claim, path in self._kwargs[
-                                                    "persistent_volume_claims"
-                                                ].items()
-                                            ]
-                                            if self._kwargs["persistent_volume_claims"]
-                                            is not None
-                                            else []
-                                        ),
-                                    )
-                                ],
-                                node_selector=self._kwargs.get("node_selector"),
-                                restart_policy="Never",
-
-                                set_hostname_as_fqdn=True, # configure pod hostname as pod's FQDN
-                                share_process_namespace=False, # default
-                                subdomain=subdomain, # FQDN = <hostname>.<subdomain>.<pod namespace>.svc.<cluster domain>
-
-                                service_account_name=self._kwargs["service_account"],
-                                termination_grace_period_seconds=0,
-                                tolerations=[
-                                    client.V1Toleration(**toleration)
-                                    for toleration in self._kwargs.get("tolerations")
-                                    or []
-                                ],
-                                volumes=(
-                                    [
-                                        client.V1Volume(
-                                            name="tmpfs-ephemeral-volume",
-                                            empty_dir=client.V1EmptyDirVolumeSource(
-                                                medium="Memory",
-                                                size_limit="{}Mi".format(tmpfs_size),
-                                            ),
-                                        )
-                                    ]
-                                    if tmpfs_enabled
-                                    else []
-                                )
-                                + (
-                                    [
-                                        client.V1Volume(
-                                            name="dhsm",
-                                            empty_dir=client.V1EmptyDirVolumeSource(
-                                                medium="Memory",
-                                                size_limit="{}Mi".format(shared_memory),
-                                            )
-                                        )
-                                    ]
-                                    if shared_memory else []
-                                )
-                                + (
-                                    [
-                                        client.V1Volume(
-                                            name=claim,
-                                            persistent_volume_claim=client.V1PersistentVolumeClaimVolumeSource(
-                                                claim_name=claim
-                                            ),
-                                        )
-                                        for claim in self._kwargs[
-                                            "persistent_volume_claims"
-                                        ].keys()
-                                    ]
-                                    if self._kwargs["persistent_volume_claims"]
-                                    is not None
-                                    else []
-                                ),
-                            ),
-                        ),
-                    ),
-                ),
-            )
-
-        if "num_parallel" in self._kwargs and self._kwargs["num_parallel"] >= 1:
-
-            try:
-                import jobset
-            except ImportError:
-                _install_jobset()
-                import jobset
-
-            main_commands = copy.copy(self._kwargs["command"])
-            main_commands[-1] = main_commands[-1].replace(
-                "[multinode-args]", "--split-index 0"
-            )
-
-            task_id = self._kwargs["attrs"]["metaflow.task_id"]
-            secondary_commands = copy.copy(self._kwargs["command"])
-            # RANK needs +1 because control node is not in the worker index group, yet we want global nodes.
-            # Technically, control and worker could be same replicated job type, but cleaner to separate for future use cases.
-            secondary_commands[-1] = secondary_commands[-1].replace(
-                "[multinode-args]", "--split-index `expr $RANK + 1`"
-            )
-            secondary_commands[-1] = secondary_commands[-1].replace(
-                "ubf_control", "ubf_task"
-            )
-            secondary_commands[-1] = secondary_commands[-1].replace(
-                task_id,
-                task_id.replace("control-", "") + "-node-`expr $RANK + 1`",
-            )
-
-            if passwordless_ssh:
-                if not os.path.exists("/usr/sbin/sshd"):
-                    raise KubernetesJobException(
-                        "This @parallel decorator requires sshd to be installed in the container image."
-                        "Please install OpenSSH."
-                    )
-
-                # run sshd in background
-                main_commands[-1] = "/usr/sbin/sshd -D & %s" % main_commands[-1]
-                secondary_commands[-1] = "/usr/sbin/sshd -D & %s" % secondary_commands[-1]
-
-            replicated_jobs = [_get_replicated_job("control", 1, main_commands)]
-            if self._kwargs["num_parallel"] > 1:
-                replicated_jobs.append(
-                    _get_replicated_job("worker", self._kwargs["num_parallel"] - 1, secondary_commands)
-                )
-
-            self._jobset = jobset.models.jobset_v1alpha2_job_set.JobsetV1alpha2JobSet(
-                api_version="jobset.x-k8s.io/v1alpha2",
-                kind="JobSet",
+        return client.V1JobSpec(
+            # Retries are handled by Metaflow when it is responsible for
+            # executing the flow. The responsibility is moved to Kubernetes
+            # when Argo Workflows is responsible for the execution.
+            backoff_limit=self._kwargs.get("retries", 0),
+            completions=self._kwargs.get("completions", 1),
+            ttl_seconds_after_finished=7
+            * 60
+            * 60  # Remove job after a week. TODO: Make this configurable
+            * 24,
+            template=client.V1PodTemplateSpec(
                 metadata=client.V1ObjectMeta(
                     annotations=self._kwargs.get("annotations", {}),
                     labels=self._kwargs.get("labels", {}),
-                    name=jobset_name,
                     namespace=self._kwargs["namespace"],
                 ),
-                spec=jobset.models.jobset_v1alpha2_job_set_spec.JobsetV1alpha2JobSetSpec(
-                    network=jobset.models.jobset_v1alpha2_network.JobsetV1alpha2Network(
-                        enable_dns_hostnames=True if not self._kwargs['attrs']['requires_passwordless_ssh'] else False,
-                        subdomain=subdomain
-                    ),
-                    replicated_jobs=replicated_jobs
-                ),
-            )
-            self._passwordless_ssh_service = _get_passwordless_ssh_service()
-        else:
-            self._job = client.V1Job(
-                api_version="batch/v1",
-                kind="Job",
-                metadata=client.V1ObjectMeta(
-                    # Annotations are for humans
-                    annotations=self._kwargs.get("annotations", {}),
-                    # While labels are for Kubernetes
-                    labels=self._kwargs.get("labels", {}),
-                    generate_name=self._kwargs["generate_name"],
-                    namespace=self._kwargs["namespace"],  # Defaults to `default`
-                ),
-                spec=client.V1JobSpec(
-                    # Retries are handled by Metaflow when it is responsible for
-                    # executing the flow. The responsibility is moved to Kubernetes
-                    # when Argo Workflows is responsible for the execution.
-                    backoff_limit=self._kwargs.get("retries", 0),
-                    completions=1,  # A single non-indexed pod job
-                    ttl_seconds_after_finished=7
-                    * 60
-                    * 60  # Remove job after a week. TODO: Make this configurable
-                    * 24,
-                    template=client.V1PodTemplateSpec(
-                        metadata=client.V1ObjectMeta(
-                            annotations=self._kwargs.get("annotations", {}),
-                            labels=self._kwargs.get("labels", {}),
-                            namespace=self._kwargs["namespace"],
-                        ),
-                        spec=client.V1PodSpec(
-                            # Timeout is set on the pod and not the job (important!)
-                            active_deadline_seconds=self._kwargs["timeout_in_seconds"],
-                            # TODO (savin): Enable affinities for GPU scheduling.
-                            # affinity=?,
-                            containers=[
-                                client.V1Container(
-                                    command=self._kwargs["command"],
-                                    ports=[
-                                        client.V1ContainerPort(
-                                            container_port=int(self._kwargs["port"])
-                                        )
-                                    ]
-                                    if "port" in self._kwargs and self._kwargs["port"]
-                                    else None,
-                                    env=[
-                                        client.V1EnvVar(name=k, value=str(v))
-                                        for k, v in self._kwargs.get(
-                                            "environment_variables", {}
-                                        ).items()
-                                    ]
-                                    # And some downward API magic. Add (key, value)
-                                    # pairs below to make pod metadata available
-                                    # within Kubernetes container.
-                                    + [
-                                        client.V1EnvVar(
-                                            name=k,
-                                            value_from=client.V1EnvVarSource(
-                                                field_ref=client.V1ObjectFieldSelector(
-                                                    field_path=str(v)
-                                                )
-                                            ),
-                                        )
-                                        for k, v in {
-                                            "METAFLOW_KUBERNETES_POD_NAMESPACE": "metadata.namespace",
-                                            "METAFLOW_KUBERNETES_POD_NAME": "metadata.name",
-                                            "METAFLOW_KUBERNETES_POD_ID": "metadata.uid",
-                                            "METAFLOW_KUBERNETES_SERVICE_ACCOUNT_NAME": "spec.serviceAccountName",
-                                            "METAFLOW_KUBERNETES_NODE_IP": "status.hostIP",
-                                        }.items()
-                                    ],
-                                    env_from=[
-                                        client.V1EnvFromSource(
-                                            secret_ref=client.V1SecretEnvSource(
-                                                name=str(k),
-                                                # optional=True
-                                            )
+                spec=client.V1PodSpec(
+                    # Timeout is set on the pod and not the job (important!)
+                    active_deadline_seconds=self._kwargs["timeout_in_seconds"],
+                    # TODO (savin): Enable affinities for GPU scheduling.
+                    # affinity=?,
+                    containers=[
+                        client.V1Container(
+                            command=self._kwargs["command"],
+                            ports=[]
+                            if self._kwargs["port"] is None
+                            else [
+                                client.V1ContainerPort(
+                                    container_port=int(self._kwargs["port"])
+                                )
+                            ],
+                            env=[
+                                client.V1EnvVar(name=k, value=str(v))
+                                for k, v in self._kwargs.get(
+                                    "environment_variables", {}
+                                ).items()
+                            ]
+                            # And some downward API magic. Add (key, value)
+                            # pairs below to make pod metadata available
+                            # within Kubernetes container.
+                            + [
+                                client.V1EnvVar(
+                                    name=k,
+                                    value_from=client.V1EnvVarSource(
+                                        field_ref=client.V1ObjectFieldSelector(
+                                            field_path=str(v)
                                         )
-                                        for k in list(self._kwargs.get("secrets", []))
-                                        + KUBERNETES_SECRETS.split(",")
-                                        if k
-                                    ],
-                                    image=self._kwargs["image"],
-                                    image_pull_policy=self._kwargs["image_pull_policy"],
-                                    name=self._kwargs["step_name"].replace("_", "-"),
-                                    resources=client.V1ResourceRequirements(
-                                        requests={
-                                            "cpu": str(self._kwargs["cpu"]),
-                                            "memory": "%sM"
-                                            % str(self._kwargs["memory"]),
-                                            "ephemeral-storage": "%sM"
-                                            % str(self._kwargs["disk"]),
-                                        },
-                                        limits={
-                                            "%s.com/gpu".lower()
-                                            % self._kwargs["gpu_vendor"]: str(
-                                                self._kwargs["gpu"]
-                                            )
-                                            for k in [0]
-                                            # Don't set GPU limits if gpu isn't specified.
-                                            if self._kwargs["gpu"] is not None
-                                        },
-                                    ),
-                                    volume_mounts=(
-                                        [
-                                            client.V1VolumeMount(
-                                                mount_path=self._kwargs.get(
-                                                    "tmpfs_path"
-                                                ),
-                                                name="tmpfs-ephemeral-volume",
-                                            )
-                                        ]
-                                        if tmpfs_enabled
-                                        else []
-                                    )
-                                    + (
-                                        [
-                                            client.V1VolumeMount(
-                                                mount_path=path, name=claim
-                                            )
-                                            for claim, path in self._kwargs[
-                                                "persistent_volume_claims"
-                                            ].items()
-                                        ]
-                                        if self._kwargs["persistent_volume_claims"]
-                                        is not None
-                                        else []
                                     ),
                                 )
+                                for k, v in {
+                                    "METAFLOW_KUBERNETES_POD_NAMESPACE": "metadata.namespace",
+                                    "METAFLOW_KUBERNETES_POD_NAME": "metadata.name",
+                                    "METAFLOW_KUBERNETES_POD_ID": "metadata.uid",
+                                    "METAFLOW_KUBERNETES_SERVICE_ACCOUNT_NAME": "spec.serviceAccountName",
+                                    "METAFLOW_KUBERNETES_NODE_IP": "status.hostIP",
+                                }.items()
+                            ]
+                            + [
+                                client.V1EnvVar(name=k, value=str(v))
+                                for k, v in inject_tracing_vars({}).items()
                             ],
-                            node_selector=self._kwargs.get("node_selector"),
-                            # TODO (savin): Support image_pull_secrets
-                            # image_pull_secrets=?,
-                            # TODO (savin): Support preemption policies
-                            # preemption_policy=?,
-                            #
-                            # A Container in a Pod may fail for a number of
-                            # reasons, such as because the process in it exited
-                            # with a non-zero exit code, or the Container was
-                            # killed due to OOM etc. If this happens, fail the pod
-                            # and let Metaflow handle the retries.
-                            restart_policy="Never",
-                            service_account_name=self._kwargs["service_account"],
-                            # Terminate the container immediately on SIGTERM
-                            termination_grace_period_seconds=0,
-                            tolerations=[
-                                client.V1Toleration(**toleration)
-                                for toleration in self._kwargs.get("tolerations") or []
+                            env_from=[
+                                client.V1EnvFromSource(
+                                    secret_ref=client.V1SecretEnvSource(
+                                        name=str(k),
+                                        # optional=True
+                                    )
+                                )
+                                for k in list(self._kwargs.get("secrets", []))
+                                + KUBERNETES_SECRETS.split(",")
+                                if k
                             ],
-                            volumes=(
+                            image=self._kwargs["image"],
+                            image_pull_policy=self._kwargs["image_pull_policy"],
+                            name=self._kwargs["step_name"].replace("_", "-"),
+                            resources=client.V1ResourceRequirements(
+                                requests={
+                                    "cpu": str(self._kwargs["cpu"]),
+                                    "memory": "%sM" % str(self._kwargs["memory"]),
+                                    "ephemeral-storage": "%sM"
+                                    % str(self._kwargs["disk"]),
+                                },
+                                limits={
+                                    "%s.com/gpu".lower()
+                                    % self._kwargs["gpu_vendor"]: str(
+                                        self._kwargs["gpu"]
+                                    )
+                                    for k in [0]
+                                    # Don't set GPU limits if gpu isn't specified.
+                                    if self._kwargs["gpu"] is not None
+                                },
+                            ),
+                            volume_mounts=(
                                 [
-                                    client.V1Volume(
+                                    client.V1VolumeMount(
+                                        mount_path=self._kwargs.get("tmpfs_path"),
                                         name="tmpfs-ephemeral-volume",
-                                        empty_dir=client.V1EmptyDirVolumeSource(
-                                            medium="Memory",
-                                            # Add default unit as ours differs from Kubernetes default.
-                                            size_limit="{}Mi".format(tmpfs_size),
-                                        ),
                                     )
                                 ]
                                 if tmpfs_enabled
@@ -610,24 +177,122 @@ class KubernetesJob(object):
                             )
                             + (
                                 [
-                                    client.V1Volume(
-                                        name=claim,
-                                        persistent_volume_claim=client.V1PersistentVolumeClaimVolumeSource(
-                                            claim_name=claim
-                                        ),
+                                    client.V1VolumeMount(
+                                        mount_path="/dev/shm", name="dhsm"
                                     )
-                                    for claim in self._kwargs[
+                                ]
+                                if shared_memory
+                                else []
+                            )
+                            + (
+                                [
+                                    client.V1VolumeMount(mount_path=path, name=claim)
+                                    for claim, path in self._kwargs[
                                         "persistent_volume_claims"
-                                    ].keys()
+                                    ].items()
                                 ]
                                 if self._kwargs["persistent_volume_claims"] is not None
                                 else []
                             ),
-                            # TODO (savin): Set termination_message_policy
-                        ),
+                        )
+                    ],
+                    node_selector=self._kwargs.get("node_selector"),
+                    # TODO (savin): Support image_pull_secrets
+                    # image_pull_secrets=?,
+                    # TODO (savin): Support preemption policies
+                    # preemption_policy=?,
+                    #
+                    # A Container in a Pod may fail for a number of
+                    # reasons, such as because the process in it exited
+                    # with a non-zero exit code, or the Container was
+                    # killed due to OOM etc. If this happens, fail the pod
+                    # and let Metaflow handle the retries.
+                    restart_policy="Never",
+                    service_account_name=self._kwargs["service_account"],
+                    # Terminate the container immediately on SIGTERM
+                    termination_grace_period_seconds=0,
+                    tolerations=[
+                        client.V1Toleration(**toleration)
+                        for toleration in self._kwargs.get("tolerations") or []
+                    ],
+                    volumes=(
+                        [
+                            client.V1Volume(
+                                name="tmpfs-ephemeral-volume",
+                                empty_dir=client.V1EmptyDirVolumeSource(
+                                    medium="Memory",
+                                    # Add default unit as ours differs from Kubernetes default.
+                                    size_limit="{}Mi".format(tmpfs_size),
+                                ),
+                            )
+                        ]
+                        if tmpfs_enabled
+                        else []
+                    )
+                    + (
+                        [
+                            client.V1Volume(
+                                name="dhsm",
+                                empty_dir=client.V1EmptyDirVolumeSource(
+                                    medium="Memory",
+                                    size_limit="{}Mi".format(shared_memory),
+                                ),
+                            )
+                        ]
+                        if shared_memory
+                        else []
+                    )
+                    + (
+                        [
+                            client.V1Volume(
+                                name=claim,
+                                persistent_volume_claim=client.V1PersistentVolumeClaimVolumeSource(
+                                    claim_name=claim
+                                ),
+                            )
+                            for claim in self._kwargs["persistent_volume_claims"].keys()
+                        ]
+                        if self._kwargs["persistent_volume_claims"] is not None
+                        else []
                     ),
+                    # TODO (savin): Set termination_message_policy
                 ),
-            )
+            ),
+        )
+
+    def create(self):
+        # A discerning eye would notice and question the choice of using the
+        # V1Job construct over the V1Pod construct given that we don't rely much
+        # on any of the V1Job semantics. The major reasons at the moment are -
+        #     1. It makes the Kubernetes UIs (Octant, Lens) a bit easier on
+        #        the eyes, although even that can be questioned.
+        #     2. AWS Step Functions, at the moment (Apr' 22) only supports
+        #        executing Jobs and not Pods as part of it's publicly declared
+        #        API. When we ship the AWS Step Functions integration with EKS,
+        #        it will hopefully lessen our workload.
+        #
+        # Note: This implementation ensures that there is only one unique Pod
+        # (unique UID) per Metaflow task attempt.
+        client = self._client.get()
+
+        # tmpfs variables
+        use_tmpfs = self._kwargs["use_tmpfs"]
+        tmpfs_size = self._kwargs["tmpfs_size"]
+        tmpfs_enabled = use_tmpfs or (tmpfs_size and not use_tmpfs)
+
+        self._job = client.V1Job(
+            api_version="batch/v1",
+            kind="Job",
+            metadata=client.V1ObjectMeta(
+                # Annotations are for humans
+                annotations=self._kwargs.get("annotations", {}),
+                # While labels are for Kubernetes
+                labels=self._kwargs.get("labels", {}),
+                generate_name=self._kwargs["generate_name"],
+                namespace=self._kwargs["namespace"],  # Defaults to `default`
+            ),
+            spec=self.create_job_spec(),
+        )
         return self
 
     def execute(self):
@@ -638,53 +303,19 @@ class KubernetesJob(object):
             #       achieve the guarantees that we are seeking.
             #       https://github.com/kubernetes/enhancements/issues/1040
             #       Hopefully, we will be able to get creative with kube-batch
-
-            if "num_parallel" in self._kwargs and self._kwargs["num_parallel"] >= 1:
-                # TODO (Eddie): this is kinda gross. fix it.
-                if self._kwargs["attrs"]["requires_passwordless_ssh"]:
-                    api_instance = client.CoreV1Api()
-                    api_response = api_instance.create_namespaced_service(namespace=self._kwargs['namespace'], body=self._passwordless_ssh_service)
-
-                with client.ApiClient() as api_client:
-                    api_instance = client.CustomObjectsApi(api_client)
-
-                response = api_instance.create_namespaced_custom_object(
-                    body=self._jobset,
-                    group="jobset.x-k8s.io",
-                    version="v1alpha2",
-                    namespace=self._kwargs["namespace"],
-                    plural="jobsets",
-                )
-
-                # HACK: Give K8s some time to actually create the job
-                time.sleep(10)
-
-                # TODO (Eddie): Remove hack and make RunningJobSet.
-                # There are many jobs running that should be monitored.
-                job_name = "%s-control-0" % response["metadata"]["name"]
-                fake_id = 123
-                return RunningJob(
-                    client=self._client,
-                    name=job_name,
-                    uid=fake_id,
-                    namespace=response["metadata"]["namespace"],
-                )
-
-            else:
-                response = (
-                    client.BatchV1Api()
-                    .create_namespaced_job(
-                        body=self._job, namespace=self._kwargs["namespace"]
-                    )
-                    .to_dict()
-                )
-                return RunningJob(
-                    client=self._client,
-                    name=response["metadata"]["name"],
-                    uid=response["metadata"]["uid"],
-                    namespace=response["metadata"]["namespace"],
+            response = (
+                client.BatchV1Api()
+                .create_namespaced_job(
+                    body=self._job, namespace=self._kwargs["namespace"]
                 )
-
+                .to_dict()
+            )
+            return RunningJob(
+                client=self._client,
+                name=response["metadata"]["name"],
+                uid=response["metadata"]["uid"],
+                namespace=response["metadata"]["namespace"],
+            )
         except client.rest.ApiException as e:
             raise KubernetesJobException(
                 "Unable to launch Kubernetes job.\n %s"
@@ -740,6 +371,7 @@ class KubernetesJob(object):
 
 
 class RunningJob(object):
+
     # State Machine implementation for the lifecycle behavior documented in
     # https://kubernetes.io/docs/concepts/workloads/pods/pod-lifecycle/
     #
@@ -793,7 +425,7 @@ class RunningJob(object):
         def best_effort_kill():
             try:
                 self.kill()
-            except:
+            except Exception as ex:
                 pass
 
         atexit.register(best_effort_kill)
@@ -857,6 +489,7 @@ class RunningJob(object):
         # 3. If the pod object hasn't shown up yet, we set the parallelism to 0
         #    to preempt it.
         client = self._client.get()
+
         if not self.is_done:
             if self.is_running:
                 # Case 1.