import os
from flytekit.contrib.notebook.tasks import spark_notebook
from flytekit.sdk.tasks import inputs, outputs
from flytekit.sdk.types import Types
from flytekit.sdk.workflow import workflow_class, Input
interactive_spark = spark_notebook(
notebook_path=os.path.join(os.path.dirname(os.path.abspath(__file__)), "spark-notebook.ipynb"),
inputs=inputs(partitions=Types.Integer),
outputs=outputs(pi=Types.Float),
)
@workflow_class
class FlyteNotebookSparkWorkflow(object):
partitions = Input(Types.Integer, default=100)
out1 = interactive_spark(partitions=partitions)
def get_sdk_node(
self,
pipeline_context,
instance,
pipeline_run,
step_key,
task_type=constants.SdkTaskType.PYTHON_TASK,
cache_version="",
retries=0,
interruptible=False,
deprecated="",
storage_request=None,
cpu_request=None,
gpu_request=None,
memory_request=None,
storage_limit=None,
cpu_limit=None,
gpu_limit=None,
memory_limit=None,
cache=False,
timeout=datetime.timedelta(seconds=0),
environment=None,
):
execution_step = self.execution_plan.get_step_by_key(step_key)
flyte_inputs = self.flyte_inputs(execution_step.step_input_dict,
execution_step.solid_name)
flyte_outputs = self.flyte_outputs(execution_step.step_output_dict,
execution_step.solid_name)
def wrapper(wf_params, *args, **kwargs): # pylint: disable=unused-argument
# TODO: We can't update config values via inputs from Flyte, because they are immutable
plan = self.execution_plan.build_subset_plan([step_key])
for param, arg in kwargs.items():
self.inject_intermediates(pipeline_context, execution_step,
param, arg)
results = list(
execute_plan(
plan,
instance,
run_config=self.run_config,
pipeline_run=pipeline_run,
))
for result in results:
step_context = pipeline_context.for_step(execution_step)
self.output_value(step_context, step_key, result,
execution_step, kwargs)
# This will take the wrapper definition and re-create it with explicit parameters as keyword argumentss
wrapper = forge.sign(forge.arg("wf_params"),
*map(forge.arg, flyte_inputs.keys()),
*map(forge.arg, flyte_outputs.keys()))(wrapper)
# flytekit uses this name for an internal representation, make it unique to the step key
wrapper.__name__ = execution_step.solid_name
task = sdk_runnable.SdkRunnableTask(
task_function=wrapper,
task_type=task_type,
discovery_version=cache_version,
retries=retries,
interruptible=interruptible,
deprecated=deprecated,
storage_request=storage_request,
cpu_request=cpu_request,
gpu_request=gpu_request,
memory_request=memory_request,
storage_limit=storage_limit,
cpu_limit=cpu_limit,
gpu_limit=gpu_limit,
memory_limit=memory_limit,
discoverable=cache,
timeout=timeout,
environment=environment,
custom={},
)
if flyte_inputs:
task = inputs(task, **flyte_inputs)
if flyte_outputs:
task = outputs(task, **flyte_outputs)
return task
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from flytekit.sdk.types import Types
from flytekit.sdk.tasks import inputs, outputs
from flytekit.sdk.workflow import workflow_class, Input
from flytekit.contrib.notebook.tasks import python_notebook
# The path
interactive_python = python_notebook(
notebook_path="./notebook-task-examples/python-notebook.ipynb",
inputs=inputs(pi=Types.Float),
outputs=outputs(out=Types.Float),
cpu_request="1",
memory_request="1G")
@workflow_class
class FlyteNotebookWorkflow(object):
out2 = interactive_python(pi=3.14)