Kubernetes Overview

from metaflow import FlowSpec, step, kubernetes

class MyFlow(FlowSpec):
    
    @step
    def start(self):
        print("Running locally")
        self.next(self.train)
    
    @kubernetes(cpu=4, memory=16000, gpu=1)
    @step
    def train(self):
        # This step runs as a Kubernetes Job
        print("Training model on Kubernetes with GPU")
        self.model = train_model()
        self.next(self.end)
    
    @step
    def end(self):
        print("Back to local execution")

if __name__ == '__main__':
    MyFlow()

@kubernetes(cpu=8, memory=32000)
@step
def process(self):
    # Runs in a single pod with 8 CPUs and 32GB RAM
    result = process_data(self.data)
    self.next(self.end)

from metaflow import kubernetes, parallel

@kubernetes(cpu=4, memory=16000)
@parallel(num_nodes=4)
@step
def distributed_training(self):
    from metaflow import current
    
    # Access node information
    print(f"Node {current.parallel.node_index} of {current.parallel.num_nodes}")
    print(f"Main IP: {current.parallel.main_ip}")
    
    # Your distributed training code here
    if current.parallel.node_index == 0:
        # Control node logic
        coordinate_training()
    else:
        # Worker node logic
        run_worker()
    
    self.next(self.end)

# Basic configuration
export METAFLOW_KUBERNETES_NAMESPACE=metaflow
export METAFLOW_KUBERNETES_SERVICE_ACCOUNT=metaflow-service-account

# Container configuration
export METAFLOW_KUBERNETES_CONTAINER_IMAGE=myregistry/metaflow:latest
export METAFLOW_KUBERNETES_CONTAINER_REGISTRY=myregistry.io

# Datastore configuration (choose one)
export METAFLOW_DATASTORE_SYSROOT_S3=s3://my-metaflow-bucket
# or
export METAFLOW_DATASTORE_SYSROOT_AZURE=wasbs://container@account.blob.core.windows.net/
# or
export METAFLOW_DATASTORE_SYSROOT_GS=gs://my-metaflow-bucket

@kubernetes(
    cpu=4,              # Number of CPU cores
    memory=16000,       # Memory in MB
    disk=50000,         # Ephemeral disk in MB
    gpu=1,              # Number of GPUs
    gpu_vendor='nvidia' # GPU vendor: 'nvidia' or 'amd'
)
@step
def compute_intensive_step(self):
    # Your code here
    pass

@kubernetes(
    node_selector={
        'node.kubernetes.io/instance-type': 'g4dn.xlarge',
        'topology.kubernetes.io/zone': 'us-east-1a'
    }
)
@step
def targeted_step(self):
    pass

@kubernetes(
    tolerations=[
        {
            'key': 'dedicated',
            'operator': 'Equal',
            'value': 'ml-workloads',
            'effect': 'NoSchedule'
        }
    ]
)
@step
def tolerated_step(self):
    pass

@kubernetes(
    persistent_volume_claims={
        'data-pvc': '/mnt/data',
        'models-pvc': '/mnt/models'
    }
)
@step
def storage_step(self):
    # Access mounted volumes
    with open('/mnt/data/input.txt') as f:
        data = f.read()
    pass

@kubernetes(
    secrets=['my-db-credentials', 'api-keys']
)
@step
def secure_step(self):
    # Secrets are mounted as environment variables
    import os
    api_key = os.environ.get('API_KEY')
    pass

# List all Metaflow jobs in your namespace
kubectl get jobs -n metaflow -l app.kubernetes.io/part-of=metaflow

# View job details
kubectl describe job <job-name> -n metaflow

# View pod logs
kubectl logs -n metaflow <pod-name>

# Follow logs in real-time
kubectl logs -n metaflow <pod-name> -f

# Get pod status
kubectl get pods -n metaflow -l metaflow/flow_name=MyFlow

# Inspect failed pod
kubectl describe pod <pod-name> -n metaflow

# View events
kubectl get events -n metaflow --sort-by='.lastTimestamp'