def test_reconcile_pods(self, mock_uuid): mock_uuid.return_value = self.static_uuid path = sys.path[ 0] + '/tests/kubernetes/pod_generator_base_with_secrets.yaml' base_pod = PodGenerator(pod_template_file=path, extract_xcom=False).gen_pod() mutator_pod = k8s.V1Pod( metadata=k8s.V1ObjectMeta( name="name2", labels={"bar": "baz"}, ), spec=k8s.V1PodSpec( containers=[ k8s.V1Container( image='', name='name', command=['/bin/command2.sh', 'arg2'], volume_mounts=[ k8s.V1VolumeMount( mount_path="/foo/", name="example-kubernetes-test-volume2") ], ) ], volumes=[ k8s.V1Volume( host_path=k8s.V1HostPathVolumeSource(path="/tmp/"), name="example-kubernetes-test-volume2", ) ], ), ) result = PodGenerator.reconcile_pods(base_pod, mutator_pod) expected: k8s.V1Pod = self.expected expected.metadata.name = "name2" expected.metadata.labels['bar'] = 'baz' expected.spec.volumes = expected.spec.volumes or [] expected.spec.volumes.append( k8s.V1Volume(host_path=k8s.V1HostPathVolumeSource(path="/tmp/"), name="example-kubernetes-test-volume2")) base_container: k8s.V1Container = expected.spec.containers[0] base_container.command = ['/bin/command2.sh', 'arg2'] base_container.volume_mounts = [ k8s.V1VolumeMount(mount_path="/foo/", name="example-kubernetes-test-volume2") ] base_container.name = "name" expected.spec.containers[0] = base_container result_dict = self.k8s_client.sanitize_for_serialization(result) expected_dict = self.k8s_client.sanitize_for_serialization(expected) assert result_dict == expected_dict
def _get_volumes(self) -> List[k8s.V1Volume]: def _construct_volume(name, claim, host) -> k8s.V1Volume: volume = k8s.V1Volume(name=name) if claim: volume.persistent_volume_claim = k8s.V1PersistentVolumeClaimVolumeSource( claim_name=claim) elif host: volume.host_path = k8s.V1HostPathVolumeSource(path=host, type='') else: volume.empty_dir = {} return volume volumes = { self.dags_volume_name: _construct_volume(self.dags_volume_name, self.kube_config.dags_volume_claim, self.kube_config.dags_volume_host), self.logs_volume_name: _construct_volume(self.logs_volume_name, self.kube_config.logs_volume_claim, self.kube_config.logs_volume_host) } if self.kube_config.dags_in_image: del volumes[self.dags_volume_name] # Get the SSH key from secrets as a volume if self.kube_config.git_ssh_key_secret_name: volumes[self.git_sync_ssh_secret_volume_name] = k8s.V1Volume( name=self.git_sync_ssh_secret_volume_name, secret=k8s.V1SecretVolumeSource( secret_name=self.kube_config.git_ssh_key_secret_name, items=[ k8s.V1KeyToPath(key=self.git_ssh_key_secret_key, path='ssh', mode=0o440) ])) if self.kube_config.git_ssh_known_hosts_configmap_name: volumes[self.git_sync_ssh_known_hosts_volume_name] = k8s.V1Volume( name=self.git_sync_ssh_known_hosts_volume_name, config_map=k8s.V1ConfigMapVolumeSource( name=self.kube_config.git_ssh_known_hosts_configmap_name, default_mode=0o440)) # Mount the airflow.cfg file via a configmap the user has specified if self.kube_config.airflow_configmap: config_volume_name = 'airflow-config' volumes[config_volume_name] = k8s.V1Volume( name=config_volume_name, config_map=k8s.V1ConfigMapVolumeSource( name=self.kube_config.airflow_configmap)) return list(volumes.values())
def get_kubeconfig_volume(release: OpenshiftRelease): return k8s.V1Volume( name="kubeconfig", secret=k8s.V1SecretVolumeSource( secret_name=f"{release.get_release_name()}-kubeconfig" ) )
def test_gen_pod_extract_xcom(self, mock_uuid): mock_uuid.return_value = self.static_uuid path = sys.path[0] + '/tests/kubernetes/pod_generator_base_with_secrets.yaml' pod_generator = PodGenerator(pod_template_file=path, extract_xcom=True) result = pod_generator.gen_pod() result_dict = self.k8s_client.sanitize_for_serialization(result) container_two = { 'name': 'airflow-xcom-sidecar', 'image': "alpine", 'command': ['sh', '-c', PodDefaults.XCOM_CMD], 'volumeMounts': [{'name': 'xcom', 'mountPath': '/airflow/xcom'}], 'resources': {'requests': {'cpu': '1m'}}, } self.expected.spec.containers.append(container_two) base_container: k8s.V1Container = self.expected.spec.containers[0] base_container.volume_mounts = base_container.volume_mounts or [] base_container.volume_mounts.append(k8s.V1VolumeMount(name="xcom", mount_path="/airflow/xcom")) self.expected.spec.containers[0] = base_container self.expected.spec.volumes = self.expected.spec.volumes or [] self.expected.spec.volumes.append( k8s.V1Volume( name='xcom', empty_dir={}, ) ) result_dict = self.k8s_client.sanitize_for_serialization(result) expected_dict = self.k8s_client.sanitize_for_serialization(self.expected) assert result_dict == expected_dict
def to_volume_secret(self) -> Tuple[k8s.V1Volume, k8s.V1VolumeMount]: """Converts to volume secret""" vol_id = f'secretvol{uuid.uuid4()}' volume = k8s.V1Volume(name=vol_id, secret=k8s.V1SecretVolumeSource(secret_name=self.secret)) if self.items: volume.secret.items = self.items return (volume, k8s.V1VolumeMount(mount_path=self.deploy_target, name=vol_id, read_only=True))
def test_to_volume_secret(self, mock_uuid): mock_uuid.return_value = '0' secret = Secret('volume', '/etc/foo', 'secret_b') assert secret.to_volume_secret() == ( k8s.V1Volume(name='secretvol0', secret=k8s.V1SecretVolumeSource(secret_name='secret_b')), k8s.V1VolumeMount(mount_path='/etc/foo', name='secretvol0', read_only=True), )
class PodDefaults: """ Static defaults for Pods """ XCOM_MOUNT_PATH = '/airflow/xcom' SIDECAR_CONTAINER_NAME = 'airflow-xcom-sidecar' XCOM_CMD = 'trap "exit 0" INT; while true; do sleep 30; done;' VOLUME_MOUNT = k8s.V1VolumeMount( name='xcom', mount_path=XCOM_MOUNT_PATH ) VOLUME = k8s.V1Volume( name='xcom', empty_dir=k8s.V1EmptyDirVolumeSource() ) SIDECAR_CONTAINER = k8s.V1Container( name=SIDECAR_CONTAINER_NAME, command=['sh', '-c', XCOM_CMD], image='alpine', volume_mounts=[VOLUME_MOUNT], resources=k8s.V1ResourceRequirements( requests={ "cpu": "1m", } ), )
def get_kubeconfig_volume(release: OpenshiftRelease, task_group): prefix=f"{task_group}-" return k8s.V1Volume( name="kubeconfig", secret=k8s.V1SecretVolumeSource( secret_name=f"{release.get_release_name()}-{prefix if 'hosted' in task_group else ''}kubeconfig" ) )
def to_volume_secret(self) -> Tuple[k8s.V1Volume, k8s.V1VolumeMount]: vol_id = 'secretvol{}'.format(uuid.uuid4()) return (k8s.V1Volume( name=vol_id, secret=k8s.V1SecretVolumeSource(secret_name=self.secret)), k8s.V1VolumeMount(mount_path=self.deploy_target, name=vol_id, read_only=True))
def test_to_k8s_object(self): volume_config = {'persistentVolumeClaim': {'claimName': 'test-volume'}} volume = Volume(name='test-volume', configs=volume_config) expected_volume = k8s.V1Volume( name="test-volume", persistent_volume_claim={"claimName": "test-volume"}) result = volume.to_k8s_client_obj() self.assertEqual(result, expected_volume)
def test_only_mount_sub_secret(self, mock_uuid): mock_uuid.return_value = '0' items = [k8s.V1KeyToPath(key="my-username", path="/extra/path")] secret = Secret('volume', '/etc/foo', 'secret_b', items=items) assert secret.to_volume_secret() == ( k8s.V1Volume( name='secretvol0', secret=k8s.V1SecretVolumeSource(secret_name='secret_b', items=items) ), k8s.V1VolumeMount(mount_path='/etc/foo', name='secretvol0', read_only=True), )
def test_init_container(self): # GIVEN volume_mounts = [ k8s.V1VolumeMount(mount_path='/etc/foo', name='test-volume', sub_path=None, read_only=True) ] init_environments = [ k8s.V1EnvVar(name='key1', value='value1'), k8s.V1EnvVar(name='key2', value='value2'), ] init_container = k8s.V1Container( name="init-container", image="ubuntu:16.04", env=init_environments, volume_mounts=volume_mounts, command=["bash", "-cx"], args=["echo 10"], ) volume = k8s.V1Volume( name='test-volume', persistent_volume_claim=k8s.V1PersistentVolumeClaimVolumeSource(claim_name='test-volume'), ) expected_init_container = { 'name': 'init-container', 'image': 'ubuntu:16.04', 'command': ['bash', '-cx'], 'args': ['echo 10'], 'env': [{'name': 'key1', 'value': 'value1'}, {'name': 'key2', 'value': 'value2'}], 'volumeMounts': [{'mountPath': '/etc/foo', 'name': 'test-volume', 'readOnly': True}], } k = KubernetesPodOperator( namespace='default', image="ubuntu:16.04", cmds=["bash", "-cx"], arguments=["echo 10"], labels={"foo": "bar"}, name="test-" + str(random.randint(0, 1000000)), task_id="task" + self.get_current_task_name(), volumes=[volume], init_containers=[init_container], in_cluster=False, do_xcom_push=False, ) context = create_context(k) k.execute(context) actual_pod = self.api_client.sanitize_for_serialization(k.pod) self.expected_pod['spec']['initContainers'] = [expected_init_container] self.expected_pod['spec']['volumes'] = [ {'name': 'test-volume', 'persistentVolumeClaim': {'claimName': 'test-volume'}} ] assert self.expected_pod == actual_pod
def to_k8s_client_obj(self): from kubernetes.client import models as k8s resp = k8s.V1Volume(name=self.name) for k, v in self.configs.items(): snake_key = Volume._convert_to_snake_case(k) if hasattr(resp, snake_key): setattr(resp, snake_key, v) else: raise AttributeError( "V1Volume does not have attribute {}".format(k)) return resp
def test_to_volume_secret(self, mock_uuid): static_uuid = uuid.UUID('cf4a56d2-8101-4217-b027-2af6216feb48') mock_uuid.return_value = static_uuid secret = Secret('volume', '/etc/foo', 'secret_b') self.assertEqual( secret.to_volume_secret(), (k8s.V1Volume( name='secretvol' + str(static_uuid), secret=k8s.V1SecretVolumeSource(secret_name='secret_b')), k8s.V1VolumeMount(mount_path='/etc/foo', name='secretvol' + str(static_uuid), read_only=True)))
def to_volume_secret(secret: "Secret") -> typing.Tuple[k8s.V1Volume, k8s.V1VolumeMount]: """Converts to volume secret""" vol_id = f"secretvol{uuid.uuid4()}" volume = k8s.V1Volume( name=vol_id, secret=k8s.V1SecretVolumeSource(secret_name=secret.secret) ) # if secret.items: # volume.secret.items = self.items return ( volume, k8s.V1VolumeMount(mount_path=secret.deploy_target, name=vol_id, read_only=True), )
def _construct_volume(name, claim, host) -> k8s.V1Volume: volume = k8s.V1Volume(name=name) if claim: volume.persistent_volume_claim = k8s.V1PersistentVolumeClaimVolumeSource( claim_name=claim) elif host: volume.host_path = k8s.V1HostPathVolumeSource(path=host, type='') else: volume.empty_dir = {} return volume
def to_volume_secret(self): import kubernetes.client.models as k8s vol_id = 'secretvol{}'.format(uuid.uuid4()) if self.deploy_target: volume_mount = k8s.V1VolumeMount(mount_path=self.deploy_target, name=vol_id, read_only=True) else: volume_mount = None return (k8s.V1Volume( name=vol_id, secret=k8s.V1SecretVolumeSource(secret_name=self.secret)), volume_mount)
def to_k8s_client_obj( volume_spec: typing.Dict[str, typing.Any]) -> k8s.V1Volume: name = volume_spec.get("name") volume = k8s.V1Volume(name=name) for k, v in volume_spec.items(): snake_key = _convert_to_snake_case(k) if hasattr(volume, snake_key): setattr(volume, snake_key, v) else: raise AttributeError( "k8s.V1Volume does not have attribute {}".format(k)) return volume
def to_k8s_client_obj(self) -> k8s.V1Volume: """ Converts to k8s object. :return: Volume Mount k8s object """ resp = k8s.V1Volume(name=self.name) for k, v in self.configs.items(): snake_key = Volume._convert_to_snake_case(k) if hasattr(resp, snake_key): setattr(resp, snake_key, v) else: raise AttributeError(f"V1Volume does not have attribute {k}") return resp
def test_volume_mount(self): with mock.patch.object(PodLauncher, 'log') as mock_logger: volume_mount = k8s.V1VolumeMount(name='test-volume', mount_path='/tmp/test_volume', sub_path=None, read_only=False) volume = k8s.V1Volume( name='test-volume', persistent_volume_claim=k8s. V1PersistentVolumeClaimVolumeSource(claim_name='test-volume'), ) args = [ "echo \"retrieved from mount\" > /tmp/test_volume/test.txt " "&& cat /tmp/test_volume/test.txt" ] k = KubernetesPodOperator( namespace='default', image="ubuntu:16.04", cmds=["bash", "-cx"], arguments=args, labels={"foo": "bar"}, volume_mounts=[volume_mount], volumes=[volume], name="test-" + str(random.randint(0, 1000000)), task_id="task" + self.get_current_task_name(), in_cluster=False, do_xcom_push=False, ) context = create_context(k) k.execute(context=context) mock_logger.info.assert_any_call('retrieved from mount') actual_pod = self.api_client.sanitize_for_serialization(k.pod) self.expected_pod['spec']['containers'][0]['args'] = args self.expected_pod['spec']['containers'][0]['volumeMounts'] = [{ 'name': 'test-volume', 'mountPath': '/tmp/test_volume', 'readOnly': False }] self.expected_pod['spec']['volumes'] = [{ 'name': 'test-volume', 'persistentVolumeClaim': { 'claimName': 'test-volume' } }] self.assertEqual(self.expected_pod, actual_pod)
def test_set_volume_claim_to_existing_claim(self): helm_settings = yaml.safe_load(git_sync_existing_claim) res = render_chart( 'GIT-SYNC', helm_settings, show_only=["templates/scheduler/scheduler-deployment.yaml"]) dep: k8s.V1Deployment = render_k8s_object(res[0], k8s.V1Deployment) volume_map = {vol.name: vol for vol in dep.spec.template.spec.volumes} dag_volume = volume_map['dags'] self.assertEqual( dag_volume, k8s.V1Volume( name="dags", persistent_volume_claim=k8s. V1PersistentVolumeClaimVolumeSource(claim_name='test-claim'), ), )
class PodDefaults: """ Static defaults for the PodGenerator """ XCOM_MOUNT_PATH = '/airflow/xcom' SIDECAR_CONTAINER_NAME = 'airflow-xcom-sidecar' XCOM_CMD = """import time while True: try: time.sleep(3600) except KeyboardInterrupt: exit(0) """ VOLUME_MOUNT = k8s.V1VolumeMount(name='xcom', mount_path=XCOM_MOUNT_PATH) VOLUME = k8s.V1Volume(name='xcom', empty_dir=k8s.V1EmptyDirVolumeSource()) SIDECAR_CONTAINER = k8s.V1Container(name=SIDECAR_CONTAINER_NAME, command=['python', '-c', XCOM_CMD], image='python:3.5-alpine', volume_mounts=[VOLUME_MOUNT])
def test_from_obj(self): result = PodGenerator.from_obj( { "pod_override": k8s.V1Pod( api_version="v1", kind="Pod", metadata=k8s.V1ObjectMeta(name="foo", annotations={"test": "annotation"}), spec=k8s.V1PodSpec( containers=[ k8s.V1Container( name="base", volume_mounts=[ k8s.V1VolumeMount( mount_path="/foo/", name="example-kubernetes-test-volume" ) ], ) ], volumes=[ k8s.V1Volume( name="example-kubernetes-test-volume", host_path=k8s.V1HostPathVolumeSource(path="/tmp/"), ) ], ), ) } ) result = self.k8s_client.sanitize_for_serialization(result) assert { 'apiVersion': 'v1', 'kind': 'Pod', 'metadata': { 'name': 'foo', 'annotations': {'test': 'annotation'}, }, 'spec': { 'containers': [ { 'name': 'base', 'volumeMounts': [{'mountPath': '/foo/', 'name': 'example-kubernetes-test-volume'}], } ], 'volumes': [{'hostPath': {'path': '/tmp/'}, 'name': 'example-kubernetes-test-volume'}], }, } == result result = PodGenerator.from_obj( { "KubernetesExecutor": { "annotations": {"test": "annotation"}, "volumes": [ { "name": "example-kubernetes-test-volume", "hostPath": {"path": "/tmp/"}, }, ], "volume_mounts": [ { "mountPath": "/foo/", "name": "example-kubernetes-test-volume", }, ], } } ) result_from_pod = PodGenerator.from_obj( { "pod_override": k8s.V1Pod( metadata=k8s.V1ObjectMeta(annotations={"test": "annotation"}), spec=k8s.V1PodSpec( containers=[ k8s.V1Container( name="base", volume_mounts=[ k8s.V1VolumeMount( name="example-kubernetes-test-volume", mount_path="/foo/" ) ], ) ], volumes=[k8s.V1Volume(name="example-kubernetes-test-volume", host_path="/tmp/")], ), ) } ) result = self.k8s_client.sanitize_for_serialization(result) result_from_pod = self.k8s_client.sanitize_for_serialization(result_from_pod) expected_from_pod = { 'metadata': {'annotations': {'test': 'annotation'}}, 'spec': { 'containers': [ { 'name': 'base', 'volumeMounts': [{'mountPath': '/foo/', 'name': 'example-kubernetes-test-volume'}], } ], 'volumes': [{'hostPath': '/tmp/', 'name': 'example-kubernetes-test-volume'}], }, } assert ( result_from_pod == expected_from_pod ), "There was a discrepency between KubernetesExecutor and pod_override" assert { 'apiVersion': 'v1', 'kind': 'Pod', 'metadata': { 'annotations': {'test': 'annotation'}, }, 'spec': { 'containers': [ { 'args': [], 'command': [], 'env': [], 'envFrom': [], 'name': 'base', 'ports': [], 'volumeMounts': [{'mountPath': '/foo/', 'name': 'example-kubernetes-test-volume'}], } ], 'hostNetwork': False, 'imagePullSecrets': [], 'volumes': [{'hostPath': {'path': '/tmp/'}, 'name': 'example-kubernetes-test-volume'}], }, } == result
secret_all_keys = Secret("env", None, "airflow-secrets-2") volume_mount = k8s.V1VolumeMount(name="test-volume", mount_path="/root/mount_file", sub_path=None, read_only=True) configmaps = [ k8s.V1EnvFromSource(config_map_ref=k8s.V1ConfigMapEnvSource( name="test-configmap-1")), k8s.V1EnvFromSource(config_map_ref=k8s.V1ConfigMapEnvSource( name="test-configmap-2")), ] volume = k8s.V1Volume( name="test-volume", persistent_volume_claim=k8s.V1PersistentVolumeClaimVolumeSource( claim_name="test-volume"), ) port = k8s.V1ContainerPort(name="http", container_port=80) init_container_volume_mounts = [ k8s.V1VolumeMount(mount_path="/etc/foo", name="test-volume", sub_path=None, read_only=True) ] init_environments = [ k8s.V1EnvVar(name="key1", value="value1"), k8s.V1EnvVar(name="key2", value="value2"),
executor_config_volume_mount = { "pod_override": k8s.V1Pod(spec=k8s.V1PodSpec( containers=[ k8s.V1Container( name="base", volume_mounts=[ k8s.V1VolumeMount( mount_path="/foo/", name="example-kubernetes-test-volume") ], ) ], volumes=[ k8s.V1Volume( name="example-kubernetes-test-volume", host_path=k8s.V1HostPathVolumeSource(path="/tmp/"), ) ], )), } @task(executor_config=executor_config_volume_mount) def test_volume_mount(): """ Tests whether the volume has been mounted. """ with open('/foo/volume_mount_test.txt', 'w') as foo: foo.write('Hello') return_code = os.system("cat /foo/volume_mount_test.txt") if return_code != 0:
def get_empty_dir_volume(): return k8s.V1Volume(name="tmpdir", empty_dir=k8s.V1EmptyDirVolumeSource())
def get_kubeconfig_volume(version, platform, profile): return k8s.V1Volume( name="kubeconfig", secret=k8s.V1SecretVolumeSource( secret_name=f"{version}-{platform}-{profile}-kubeconfig"))
dag = DAG( dag_id="kube-pod-operator", start_date=airflow.utils.dates.days_ago(2), schedule_interval="@daily", ) start_kube_process = BashOperator( task_id="start_kube_process", bash_command="echo upload to s3", dag=dag, ) kube_files_volume_config = {'name': 'minikube'} kube_files_volume = k8s.V1Volume(name='kube-files-volume', config_map=kube_files_volume_config) kube_files_volume_mount = k8s.V1VolumeMount(name='kube-files-volume', mount_path='/tmp/k8s') in_cluster = False kubernetes_min_pod = kubernetes_pod_operator.KubernetesPodOperator( task_id='pod-ex-minimum', name='pod-ex-minimum', cmds=['echo'], namespace='default', image='ubuntu:latest', in_cluster=in_cluster, config_file='/tmp/k8s/minikube/config', volumes=[kube_files_volume], volume_mounts=[kube_files_volume_mount]) run_another_pod = kubernetes_pod_operator.KubernetesPodOperator(
def pipeline_definition( hydrosphere_name="local", hydrosphere_address="http://hydro-serving-sidecar-serving.kubeflow.svc.cluster.local:8080", data_directory='/data/mnist', models_directory="/models/mnist", learning_rate="0.01", learning_steps="5000", batch_size="256", warmpup_count="100", model_name="mnist", application_name="mnist-app", signature_name="predict", acceptable_accuracy="0.90", ): data_pvc = k8s.V1PersistentVolumeClaimVolumeSource(claim_name="data") models_pvc = k8s.V1PersistentVolumeClaimVolumeSource(claim_name="models") data_volume = k8s.V1Volume(name="data", persistent_volume_claim=data_pvc) models_volume = k8s.V1Volume(name="models", persistent_volume_claim=models_pvc) data_volume_mount = k8s.V1VolumeMount( mount_path="{{workflow.parameters.data-directory}}", name="data") models_volume_mount = k8s.V1VolumeMount( mount_path="{{workflow.parameters.models-directory}}", name="models") hydrosphere_address_env = k8s.V1EnvVar( name="CLUSTER_ADDRESS", value="{{workflow.parameters.hydrosphere-address}}") hydrosphere_name_env = k8s.V1EnvVar( name="CLUSTER_NAME", value="{{workflow.parameters.hydrosphere-name}}") data_directory_env = k8s.V1EnvVar( name="MNIST_DATA_DIR", value="{{workflow.parameters.data-directory}}") models_directory_env = k8s.V1EnvVar( name="MNIST_MODELS_DIR", value="{{workflow.parameters.models-directory}}") model_name_env = k8s.V1EnvVar(name="MODEL_NAME", value="{{workflow.parameters.model-name}}") application_name_env = k8s.V1EnvVar( name="APPLICATION_NAME", value="{{workflow.parameters.application-name}}") signature_name_env = k8s.V1EnvVar( name="SIGNATURE_NAME", value="{{workflow.parameters.signature-name}}") acceptable_accuracy_env = k8s.V1EnvVar( name="ACCEPTABLE_ACCURACY", value="{{workflow.parameters.acceptable-accuracy}}") learning_rate_env = k8s.V1EnvVar( name="LEARNING_RATE", value="{{workflow.parameters.learning-rate}}") learning_steps_env = k8s.V1EnvVar( name="LEARNING_STEPS", value="{{workflow.parameters.learning-steps}}") batch_size_env = k8s.V1EnvVar(name="BATCH_SIZE", value="{{workflow.parameters.batch-size}}") warmup_count_env = k8s.V1EnvVar( name="WARMUP_IMAGES_AMOUNT", value="{{workflow.parameters.warmpup-count}}") # 1. Download MNIST data download = dsl.ContainerOp( name="download", image="tidylobster/mnist-pipeline-download:latest") download.add_volume(data_volume) download.add_volume_mount(data_volume_mount) download.add_env_variable(data_directory_env) # 2. Train and save a MNIST classifier using Tensorflow train = dsl.ContainerOp(name="train", image="tidylobster/mnist-pipeline-train:latest") train.after(download) train.set_memory_request('2G') train.set_cpu_request('1') train.add_volume(data_volume) train.add_volume(models_volume) train.add_volume_mount(data_volume_mount) train.add_volume_mount(models_volume_mount) train.add_env_variable(data_directory_env) train.add_env_variable(models_directory_env) train.add_env_variable(learning_rate_env) train.add_env_variable(learning_steps_env) train.add_env_variable(batch_size_env) # 3. Upload trained model to the cluster upload = dsl.ContainerOp( name="upload", image="tidylobster/mnist-pipeline-upload:latest", file_outputs={"model_version": "/model_version.txt"}) upload.after(train) upload.add_volume(models_volume) upload.add_volume_mount(models_volume_mount) upload.add_env_variable(models_directory_env) upload.add_env_variable(model_name_env) upload.add_env_variable(hydrosphere_name_env) upload.add_env_variable(hydrosphere_address_env) # 4. Deploy application deploy = dsl.ContainerOp(name="deploy", image="tidylobster/mnist-pipeline-deploy:latest", arguments=[upload.outputs["model_version"]]) deploy.after(upload) deploy.add_env_variable(hydrosphere_name_env) deploy.add_env_variable(hydrosphere_address_env) deploy.add_env_variable(application_name_env) deploy.add_env_variable(model_name_env) # 5. Test the model test = dsl.ContainerOp(name="test", image="tidylobster/mnist-pipeline-test:latest") test.after(deploy) test.add_volume(data_volume) test.add_volume_mount(data_volume_mount) test.add_env_variable(data_directory_env) test.add_env_variable(hydrosphere_address_env) test.add_env_variable(application_name_env) test.add_env_variable(signature_name_env) test.add_env_variable(warmup_count_env) test.add_env_variable(acceptable_accuracy_env) # 6. Clean environment clean = dsl.ContainerOp(name="clean", image="tidylobster/mnist-pipeline-clean:latest") clean.after(test) clean.add_volume(data_volume) clean.add_volume_mount(data_volume_mount) clean.add_env_variable(data_directory_env) clean.add_volume(models_volume) clean.add_volume_mount(models_volume_mount) clean.add_env_variable(models_directory_env)
dag_id='ai_training_run', default_args=ai_training_run_dag_default_args, schedule_interval=None, start_date=days_ago(2), tags=['training'] ) # Define Kubernetes namespace to execute DAG in namespace = 'airflow' ## Define volume details (change values as necessary to match your environment) # Dataset volume dataset_volume_pvc_existing = 'dataset-vol' dataset_volume = k8s.V1Volume( name=dataset_volume_pvc_existing, persistent_volume_claim=k8s.V1PersistentVolumeClaimVolumeSource(claim_name=dataset_volume_pvc_existing), ) dataset_volume_mount = k8s.V1VolumeMount( name=dataset_volume_pvc_existing, mount_path='/mnt/dataset', sub_path=None, read_only=False ) # Model volume model_volume_pvc_existing = 'airflow-model-vol' model_volume = k8s.V1Volume( name=model_volume_pvc_existing, persistent_volume_claim=k8s.V1PersistentVolumeClaimVolumeSource(claim_name=model_volume_pvc_existing), ) model_volume_mount = k8s.V1VolumeMount(