This is a sample for Kubeflow TFJob SDK kubeflow-tfjob.
The notebook shows how to use Kubeflow TFJob SDK to create, get, wait, check and delete tfjob.
In [1]:
from kubernetes.client import V1PodTemplateSpec
from kubernetes.client import V1ObjectMeta
from kubernetes.client import V1PodSpec
from kubernetes.client import V1Container
from kubeflow.tfjob import constants
from kubeflow.tfjob import utils
from kubeflow.tfjob import V1ReplicaSpec
from kubeflow.tfjob import V1TFJob
from kubeflow.tfjob import V1TFJobSpec
from kubeflow.tfjob import TFJobClient
Define namespace where tfjob needs to be created to. If not specified, below function defines namespace to the current one where SDK is running in the cluster, otherwise it will deploy to default namespace.
In [2]:
namespace = utils.get_default_target_namespace()
The demo only creates a worker of TFJob to run mnist sample.
In [3]:
container = V1Container(
name="tensorflow",
image="gcr.io/kubeflow-ci/tf-mnist-with-summaries:1.0",
command=[
"python",
"/var/tf_mnist/mnist_with_summaries.py",
"--log_dir=/train/logs", "--learning_rate=0.01",
"--batch_size=150"
]
)
worker = V1ReplicaSpec(
replicas=2,
restart_policy="Never",
template=V1PodTemplateSpec(
spec=V1PodSpec(
containers=[container]
)
)
)
chief = V1ReplicaSpec(
replicas=1,
restart_policy="Never",
template=V1PodTemplateSpec(
spec=V1PodSpec(
containers=[container]
)
)
)
ps = V1ReplicaSpec(
replicas=1,
restart_policy="Never",
template=V1PodTemplateSpec(
spec=V1PodSpec(
containers=[container]
)
)
)
tfjob = V1TFJob(
api_version="kubeflow.org/v1",
kind="TFJob",
metadata=V1ObjectMeta(name="mnist",namespace=namespace),
spec=V1TFJobSpec(
clean_pod_policy="None",
tf_replica_specs={"Worker": worker,
"Chief": chief,
"PS": ps}
)
)
In [4]:
tfjob_client = TFJobClient()
tfjob_client.create(tfjob, namespace=namespace)
Out[4]:
In [5]:
tfjob_client.get('mnist', namespace=namespace)
Out[5]:
In [6]:
tfjob_client.get_job_status('mnist', namespace=namespace)
Out[6]:
In [7]:
tfjob_client.wait_for_job('mnist', namespace=namespace, watch=True)
In [8]:
tfjob_client.is_job_succeeded('mnist', namespace=namespace)
Out[8]:
In [9]:
tfjob_client.get_logs('mnist', namespace=namespace)
In [10]:
tfjob_client.delete('mnist', namespace=namespace)
Out[10]:
In [ ]: