def __init__(
self,
local_path: os.PathLike = None,
remote_path: os.PathLike = None,
supported_mode: SchemaOpenMode = SchemaOpenMode.WRITE,
downloader: typing.Callable[[str, os.PathLike], None] = None,
):
if supported_mode == SchemaOpenMode.READ and remote_path is None:
raise ValueError(
"To create a FlyteSchema in read mode, remote_path is required"
)
if (supported_mode == SchemaOpenMode.WRITE and local_path is None
and FlyteContextManager.current_context().file_access is None):
raise ValueError(
"To create a FlyteSchema in write mode, local_path is required"
)
local_path = local_path or FlyteContextManager.current_context(
).file_access.get_random_local_directory()
self._local_path = local_path
# Make this field public, so that the dataclass transformer can set a value for it
# https://github.com/flyteorg/flytekit/blob/bcc8541bd6227b532f8462563fe8aac902242b21/flytekit/core/type_engine.py#L298
self.remote_path = remote_path or FlyteContextManager.current_context(
).file_access.get_random_remote_path()
self._supported_mode = supported_mode
# This is a special attribute that indicates if the data was either downloaded or uploaded
self._downloaded = False
self._downloader = downloader
def __init__(
self,
local_path: os.PathLike = None,
remote_path: str = None,
supported_mode: SchemaOpenMode = SchemaOpenMode.WRITE,
downloader: typing.Callable[[str, os.PathLike], None] = None,
):
if supported_mode == SchemaOpenMode.READ and remote_path is None:
raise ValueError(
"To create a FlyteSchema in read mode, remote_path is required"
)
if (supported_mode == SchemaOpenMode.WRITE and local_path is None
and FlyteContextManager.current_context().file_access is None):
raise ValueError(
"To create a FlyteSchema in write mode, local_path is required"
)
if local_path is None:
local_path = FlyteContextManager.current_context(
).file_access.get_random_local_directory()
self._local_path = local_path
self._remote_path = remote_path
self._supported_mode = supported_mode
# This is a special attribute that indicates if the data was either downloaded or uploaded
self._downloaded = False
self._downloader = downloader
def test_two(two_sample_inputs):
my_input = two_sample_inputs[0]
my_input_2 = two_sample_inputs[1]
@dynamic
def dt1(a: List[MyInput]) -> List[FlyteFile]:
x = []
for aa in a:
x.append(aa.main_product)
return x
with FlyteContextManager.with_context(
FlyteContextManager.current_context().with_serialization_settings(
SerializationSettings(
project="test_proj",
domain="test_domain",
version="abc",
image_config=ImageConfig(Image(name="name", fqn="image", tag="name")),
env={},
)
)
) as ctx:
with FlyteContextManager.with_context(
ctx.with_execution_state(
ctx.execution_state.with_params(
mode=ExecutionState.Mode.TASK_EXECUTION,
)
)
) as ctx:
input_literal_map = TypeEngine.dict_to_literal_map(
ctx, d={"a": [my_input, my_input_2]}, type_hints={"a": List[MyInput]}
)
dynamic_job_spec = dt1.dispatch_execute(ctx, input_literal_map)
assert len(dynamic_job_spec.literals["o0"].collection.literals) == 2
def dynamic_execute(self, task_function: Callable, **kwargs) -> Any:
"""
By the time this function is invoked, the _local_execute function should have unwrapped the Promises and Flyte
literal wrappers so that the kwargs we are working with here are now Python native literal values. This
function is also expected to return Python native literal values.
Since the user code within a dynamic task constitute a workflow, we have to first compile the workflow, and
then execute that workflow.
When running for real in production, the task would stop after the compilation step, and then create a file
representing that newly generated workflow, instead of executing it.
"""
ctx = FlyteContextManager.current_context()
if ctx.execution_state and ctx.execution_state.mode == ExecutionState.Mode.LOCAL_WORKFLOW_EXECUTION:
updated_exec_state = ctx.execution_state.with_params(
mode=ExecutionState.Mode.TASK_EXECUTION)
with FlyteContextManager.with_context(
ctx.with_execution_state(updated_exec_state)):
logger.info("Executing Dynamic workflow, using raw inputs")
return task_function(**kwargs)
if ctx.execution_state and ctx.execution_state.mode == ExecutionState.Mode.TASK_EXECUTION:
is_fast_execution = bool(
ctx.execution_state and ctx.execution_state.additional_context
and ctx.execution_state.additional_context.get(
"dynamic_addl_distro"))
if is_fast_execution:
ctx = ctx.with_serialization_settings(
SerializationSettings.new_builder(
).with_fast_serialization_settings(
FastSerializationSettings(enabled=True)).build())
return self.compile_into_workflow(ctx, task_function, **kwargs)
def test_to_python_value_without_incoming_columns():
# make a literal with a type with no columns
ctx = FlyteContextManager.current_context()
lt = TypeEngine.to_literal_type(pd.DataFrame)
df = generate_pandas()
fdt = StructuredDatasetTransformerEngine()
lit = fdt.to_literal(ctx, df, python_type=pd.DataFrame, expected=lt)
assert len(lit.scalar.structured_dataset.metadata.structured_dataset_type.
columns) == 0
# declare a new type that only has one column
# get the dataframe, make sure it has the column that was asked for.
subset_sd_type = Annotated[StructuredDataset, kwtypes(age=int)]
sd = fdt.to_python_value(ctx, lit, subset_sd_type)
assert sd.metadata.structured_dataset_type.columns[0].name == "age"
sub_df = sd.open(pd.DataFrame).all()
assert sub_df.shape[1] == 1
# check when columns are not specified, should pull both and add column information.
# todo: see the todos in the open_as, and iter_as functions in StructuredDatasetTransformerEngine
# we have to recreate the literal because the test case above filled in the metadata
lit = fdt.to_literal(ctx, df, python_type=pd.DataFrame, expected=lt)
sd = fdt.to_python_value(ctx, lit, StructuredDataset)
assert sd.metadata.structured_dataset_type.columns == []
sub_df = sd.open(pd.DataFrame).all()
assert sub_df.shape[1] == 2
# should also work if subset type is just an annotated pd.DataFrame
lit = fdt.to_literal(ctx, df, python_type=pd.DataFrame, expected=lt)
subset_pd_type = Annotated[pd.DataFrame, kwtypes(age=int)]
sub_df = fdt.to_python_value(ctx, lit, subset_pd_type)
assert sub_df.shape[1] == 1
def test_to_python_value_with_incoming_columns():
# make a literal with a type that has two columns
original_type = Annotated[pd.DataFrame, kwtypes(name=str, age=int)]
ctx = FlyteContextManager.current_context()
lt = TypeEngine.to_literal_type(original_type)
df = generate_pandas()
fdt = StructuredDatasetTransformerEngine()
lit = fdt.to_literal(ctx, df, python_type=original_type, expected=lt)
assert len(lit.scalar.structured_dataset.metadata.structured_dataset_type.
columns) == 2
# declare a new type that only has one column
# get the dataframe, make sure it has the column that was asked for.
subset_sd_type = Annotated[StructuredDataset, kwtypes(age=int)]
sd = fdt.to_python_value(ctx, lit, subset_sd_type)
assert sd.metadata.structured_dataset_type.columns[0].name == "age"
sub_df = sd.open(pd.DataFrame).all()
assert sub_df.shape[1] == 1
# check when columns are not specified, should pull both and add column information.
sd = fdt.to_python_value(ctx, lit, StructuredDataset)
assert len(sd.metadata.structured_dataset_type.columns) == 2
# should also work if subset type is just an annotated pd.DataFrame
subset_pd_type = Annotated[pd.DataFrame, kwtypes(age=int)]
sub_df = fdt.to_python_value(ctx, lit, subset_pd_type)
assert sub_df.shape[1] == 1
def test_fill_in_literal_type():
class TempEncoder(StructuredDatasetEncoder):
def __init__(self, fmt: str):
super().__init__(MyDF, "tmpfs://", supported_format=fmt)
def encode(
self,
ctx: FlyteContext,
structured_dataset: StructuredDataset,
structured_dataset_type: StructuredDatasetType,
) -> literals.StructuredDataset:
return literals.StructuredDataset(uri="")
StructuredDatasetTransformerEngine.register(TempEncoder("myavro"),
default_for_type=True)
lt = TypeEngine.to_literal_type(MyDF)
assert lt.structured_dataset_type.format == "myavro"
ctx = FlyteContextManager.current_context()
fdt = StructuredDatasetTransformerEngine()
sd = StructuredDataset(dataframe=42)
l = fdt.to_literal(ctx, sd, MyDF, lt)
# Test that the literal type is filled in even though the encode function above doesn't do it.
assert l.scalar.structured_dataset.metadata.structured_dataset_type.format == "myavro"
# Test that looking up encoders/decoders falls back to the "" encoder/decoder
empty_format_temp_encoder = TempEncoder("")
StructuredDatasetTransformerEngine.register(empty_format_temp_encoder,
default_for_type=False)
res = StructuredDatasetTransformerEngine.get_encoder(
MyDF, "tmpfs", "rando")
assert res is empty_format_temp_encoder
def dynamic_execute(self, task_function: Callable, **kwargs) -> Any:
"""
By the time this function is invoked, the local_execute function should have unwrapped the Promises and Flyte
literal wrappers so that the kwargs we are working with here are now Python native literal values. This
function is also expected to return Python native literal values.
Since the user code within a dynamic task constitute a workflow, we have to first compile the workflow, and
then execute that workflow.
When running for real in production, the task would stop after the compilation step, and then create a file
representing that newly generated workflow, instead of executing it.
"""
ctx = FlyteContextManager.current_context()
if ctx.execution_state and ctx.execution_state.mode == ExecutionState.Mode.LOCAL_WORKFLOW_EXECUTION:
updated_exec_state = ctx.execution_state.with_params(
mode=ExecutionState.Mode.TASK_EXECUTION)
with FlyteContextManager.with_context(
ctx.with_execution_state(updated_exec_state)):
logger.info("Executing Dynamic workflow, using raw inputs")
return exception_scopes.user_entry_point(task_function)(
**kwargs)
if ctx.execution_state and ctx.execution_state.mode == ExecutionState.Mode.TASK_EXECUTION:
return self.compile_into_workflow(ctx, task_function, **kwargs)
if ctx.execution_state and ctx.execution_state.mode == ExecutionState.Mode.LOCAL_TASK_EXECUTION:
return exception_scopes.user_entry_point(task_function)(**kwargs)
raise ValueError(
f"Invalid execution provided, execution state: {ctx.execution_state}"
)
def save(self, cp: typing.Union[Path, str, io.BufferedReader]):
# We have to lazy load, until we fix the imports
from flytekit.core.context_manager import FlyteContextManager
fa = FlyteContextManager.current_context().file_access
if isinstance(cp, (Path, str)):
if isinstance(cp, str):
cp = Path(cp)
if cp.is_dir():
fa.upload_directory(str(cp), self._checkpoint_dest)
else:
fname = cp.stem + cp.suffix
rpath = fa._default_remote.construct_path(False, False, self._checkpoint_dest, fname)
fa.upload(str(cp), rpath)
return
if not isinstance(cp, io.IOBase):
raise ValueError(f"Only a valid path or IOBase type (reader) should be provided, received {type(cp)}")
p = Path(self._td.name)
dest_cp = p.joinpath(self.TMP_DST_PATH)
with dest_cp.open("wb") as f:
f.write(cp.read())
rpath = fa._default_remote.construct_path(False, False, self._checkpoint_dest, self.TMP_DST_PATH)
fa.upload(str(dest_cp), rpath)
def serialize(
pkgs: typing.List[str],
settings: SerializationSettings,
local_source_root: typing.Optional[str] = None,
options: typing.Optional[Options] = None,
) -> typing.List[RegistrableEntity]:
"""
See :py:class:`flytekit.models.core.identifier.ResourceType` to match the trailing index in the file name with the
entity type.
:param options:
:param settings: SerializationSettings to be used
:param pkgs: Dot-delimited Python packages/subpackages to look into for serialization.
:param local_source_root: Where to start looking for the code.
"""
ctx = FlyteContextManager.current_context().with_serialization_settings(
settings)
with FlyteContextManager.with_context(ctx) as ctx:
# Scan all modules. the act of loading populates the global singleton that contains all objects
with module_loader.add_sys_path(local_source_root):
click.secho(
f"Loading packages {pkgs} under source root {local_source_root}",
fg="yellow")
module_loader.just_load_modules(pkgs=pkgs)
registrable_entities = get_registrable_entities(ctx, options=options)
click.secho(
f"Successfully serialized {len(registrable_entities)} flyte objects",
fg="green")
return registrable_entities
def inputs(self) -> Dict[str, Any]:
"""
Returns the inputs of the task execution in the standard Python format that is produced by
the type engine.
"""
from flytekit.control_plane.tasks.task import FlyteTask
if self._inputs is None:
client = _flyte_engine.get_client()
execution_data = client.get_task_execution_data(self.id)
# Inputs are returned inline unless they are too big, in which case a url blob pointing to them is returned.
input_map = _literal_models.LiteralMap({})
if bool(execution_data.full_inputs.literals):
input_map = execution_data.full_inputs
elif execution_data.inputs.bytes > 0:
with _common_utils.AutoDeletingTempDir() as tmp_dir:
tmp_name = os.path.join(tmp_dir.name, "inputs.pb")
_data_proxy.Data.get_data(execution_data.inputs.url,
tmp_name)
input_map = _literal_models.LiteralMap.from_flyte_idl(
_common_utils.load_proto_from_file(
_literals_pb2.LiteralMap, tmp_name))
task = FlyteTask.fetch(self.id.task_id.project,
self.id.task_id.domain,
self.id.task_id.name,
self.id.task_id.version)
self._inputs = TypeEngine.literal_map_to_kwargs(
ctx=FlyteContextManager.current_context(),
lm=input_map,
python_types=TypeEngine.guess_python_types(
task.interface.inputs),
)
return self._inputs
def inputs(self) -> Dict[str, Any]:
"""
Returns the inputs to the execution in the standard python format as dictated by the type engine.
"""
if self._inputs is None:
client = _flyte_engine.get_client()
execution_data = client.get_execution_data(self.id)
# Inputs are returned inline unless they are too big, in which case a url blob pointing to them is returned.
input_map: _literal_models.LiteralMap = _literal_models.LiteralMap(
{})
if bool(execution_data.full_inputs.literals):
input_map = execution_data.full_inputs
elif execution_data.inputs.bytes > 0:
with _common_utils.AutoDeletingTempDir() as tmp_dir:
tmp_name = _os.path.join(tmp_dir.name, "inputs.pb")
_data_proxy.Data.get_data(execution_data.inputs.url,
tmp_name)
input_map = _literal_models.LiteralMap.from_flyte_idl(
_common_utils.load_proto_from_file(
_literals_pb2.Literalmap, tmp_name))
lp_id = self.spec.launch_plan
workflow = _workflow.FlyteWorkflow.fetch(lp_id.project,
lp_id.domain, lp_id.name,
lp_id.version)
self._inputs = TypeEngine.literal_map_to_kwargs(
ctx=FlyteContextManager.current_context(),
lm=input_map,
python_types=TypeEngine.guess_python_types(
workflow.interface.inputs),
)
return self._inputs
def test_enum_type():
t = TypeEngine.to_literal_type(Color)
assert t is not None
assert t.enum_type is not None
assert t.enum_type.values
assert t.enum_type.values == [c.value for c in Color]
ctx = FlyteContextManager.current_context()
lv = TypeEngine.to_literal(ctx, Color.RED, Color, TypeEngine.to_literal_type(Color))
assert lv
assert lv.scalar
assert lv.scalar.primitive.string_value == "red"
v = TypeEngine.to_python_value(ctx, lv, Color)
assert v
assert v == Color.RED
v = TypeEngine.to_python_value(ctx, lv, str)
assert v
assert v == "red"
with pytest.raises(ValueError):
TypeEngine.to_python_value(ctx, Literal(scalar=Scalar(primitive=Primitive(string_value=str(Color.RED)))), Color)
with pytest.raises(ValueError):
TypeEngine.to_python_value(ctx, Literal(scalar=Scalar(primitive=Primitive(string_value="bad"))), Color)
with pytest.raises(AssertionError):
TypeEngine.to_literal_type(UnsupportedEnumValues)
def create(
cls,
name: str,
workflow: _annotated_workflow.WorkflowBase,
default_inputs: Dict[str, Any] = None,
fixed_inputs: Dict[str, Any] = None,
schedule: _schedule_model.Schedule = None,
notifications: List[_common_models.Notification] = None,
auth_role: _common_models.AuthRole = None,
) -> LaunchPlan:
ctx = FlyteContextManager.current_context()
default_inputs = default_inputs or {}
fixed_inputs = fixed_inputs or {}
# Default inputs come from two places, the original signature of the workflow function, and the default_inputs
# argument to this function. We'll take the latter as having higher precedence.
wf_signature_parameters = transform_inputs_to_parameters(
ctx, workflow.python_interface)
# Construct a new Interface object with just the default inputs given to get Parameters, maybe there's an
# easier way to do this, think about it later.
temp_inputs = {}
for k, v in default_inputs.items():
temp_inputs[k] = (workflow.python_interface.inputs[k], v)
temp_interface = Interface(inputs=temp_inputs, outputs={})
temp_signature = transform_inputs_to_parameters(ctx, temp_interface)
wf_signature_parameters._parameters.update(temp_signature.parameters)
# These are fixed inputs that cannot change at launch time. If the same argument is also in default inputs,
# it'll be taken out from defaults in the LaunchPlan constructor
fixed_literals = translate_inputs_to_literals(
ctx,
incoming_values=fixed_inputs,
flyte_interface_types=workflow.interface.inputs,
native_types=workflow.python_interface.inputs,
)
fixed_lm = _literal_models.LiteralMap(literals=fixed_literals)
lp = cls(
name=name,
workflow=workflow,
parameters=wf_signature_parameters,
fixed_inputs=fixed_lm,
schedule=schedule,
notifications=notifications,
auth_role=auth_role,
)
# This is just a convenience - we'll need the fixed inputs LiteralMap for when serializing the Launch Plan out
# to protobuf, but for local execution and such, why not save the original Python native values as well so
# we don't have to reverse it back every time.
default_inputs.update(fixed_inputs)
lp._saved_inputs = default_inputs
if name in cls.CACHE:
raise AssertionError(
f"Launch plan named {name} was already created! Make sure your names are unique."
)
cls.CACHE[name] = lp
return lp
def all(self) -> DF:
if self._dataframe_type is None:
raise ValueError(
"No dataframe type set. Use open() to set the local dataframe type you want to use."
)
ctx = FlyteContextManager.current_context()
return flyte_dataset_transformer.open_as(ctx, self.literal,
self._dataframe_type,
self.metadata)
def __init__(self, name: str):
self._name = name
self._cases: typing.List[Case] = []
self._last_case = False
self._condition = Condition(self)
ctx = FlyteContextManager.current_context()
# A new conditional section has been started, so lets push the context
FlyteContextManager.push_context(
ctx.enter_conditional_section().build())
def iter(self) -> Generator[DF, None, None]:
if self._dataframe_type is None:
raise ValueError(
"No dataframe type set. Use open() to set the local dataframe type you want to use."
)
ctx = FlyteContextManager.current_context()
return flyte_dataset_transformer.iter_as(
ctx,
self.literal,
self._dataframe_type,
updated_metadata=self.metadata)
def test_pb_guess_python_type():
artifact_tag = catalog_pb2.CatalogArtifactTag(artifact_id="artifact_1", name="artifact_name")
x = {"a": artifact_tag}
lt = TypeEngine.to_literal_type(catalog_pb2.CatalogArtifactTag)
gt = TypeEngine.guess_python_type(lt)
assert gt == catalog_pb2.CatalogArtifactTag
ctx = FlyteContextManager.current_context()
lm = TypeEngine.dict_to_literal_map(ctx, x, {"a": gt})
pv = TypeEngine.to_python_value(ctx, lm.literals["a"], gt)
assert pv == artifact_tag
def test_dc_dyn_directory(folders_and_files_setup):
proxy_c = MyProxyConfiguration(splat_data_dir="/tmp/proxy_splat", apriori_file="/opt/config/a_file")
proxy_p = MyProxyParameters(id="pp_id", job_i_step=1)
my_input_gcs = MyInput(
main_product=FlyteFile(folders_and_files_setup[0]),
apriori_config=MyAprioriConfiguration(
static_data_dir=FlyteDirectory("gs://my-bucket/one"),
external_data_dir=FlyteDirectory("gs://my-bucket/two"),
),
proxy_config=proxy_c,
proxy_params=proxy_p,
)
my_input_gcs_2 = MyInput(
main_product=FlyteFile(folders_and_files_setup[0]),
apriori_config=MyAprioriConfiguration(
static_data_dir=FlyteDirectory("gs://my-bucket/three"),
external_data_dir=FlyteDirectory("gs://my-bucket/four"),
),
proxy_config=proxy_c,
proxy_params=proxy_p,
)
@dynamic
def dt1(a: List[MyInput]) -> List[FlyteDirectory]:
x = []
for aa in a:
x.append(aa.apriori_config.external_data_dir)
return x
ctx = FlyteContextManager.current_context()
cb = (
ctx.new_builder()
.with_serialization_settings(
SerializationSettings(
project="test_proj",
domain="test_domain",
version="abc",
image_config=ImageConfig(Image(name="name", fqn="image", tag="name")),
env={},
)
)
.with_execution_state(ctx.execution_state.with_params(mode=ExecutionState.Mode.TASK_EXECUTION))
)
with FlyteContextManager.with_context(cb) as ctx:
input_literal_map = TypeEngine.dict_to_literal_map(
ctx, d={"a": [my_input_gcs, my_input_gcs_2]}, type_hints={"a": List[MyInput]}
)
dynamic_job_spec = dt1.dispatch_execute(ctx, input_literal_map)
assert dynamic_job_spec.literals["o0"].collection.literals[0].scalar.blob.uri == "gs://my-bucket/two"
assert dynamic_job_spec.literals["o0"].collection.literals[1].scalar.blob.uri == "gs://my-bucket/four"
def t1() -> FlyteDirectory:
user_ctx = FlyteContextManager.current_context().user_space_params
# Create a local directory to work with
p = os.path.join(user_ctx.working_directory, "test_wf")
if os.path.exists(p):
shutil.rmtree(p)
pathlib.Path(p).mkdir(parents=True)
for i in range(1, 6):
with open(os.path.join(p, f"{i}.txt"), "w") as fh:
fh.write(f"I'm file {i}\n")
return FlyteDirectory(p)
def get_type_for_output_var(self, k: str, v: Any) -> Optional[Type[Any]]:
"""
We override this method from flytekit.core.base_task Task because the dispatch_execute method uses this
interface to construct outputs. Each instance of an container_array task will however produce outputs
according to the underlying run_task interface and the array plugin handler will actually create a collection
from these individual outputs as the final output value.
"""
ctx = FlyteContextManager.current_context()
if ctx.execution_state is not None and ctx.execution_state.mode == ExecutionState.Mode.LOCAL_WORKFLOW_EXECUTION:
# In workflow execution mode we actually need to use the parent (mapper) task output interface.
return self._python_interface.outputs[k]
return self._run_task._python_interface.outputs[k]
def end_branch(
self
) -> Optional[Union[Condition, Promise, Tuple[Promise], VoidPromise]]:
"""
This should be invoked after every branch has been visited.
In case this is not local workflow execution then, we should check if this is the last case.
If so then return the promise, else return the condition
"""
if self._last_case:
# We have completed the conditional section, lets pop off the branch context
FlyteContextManager.pop_context()
ctx = FlyteContextManager.current_context()
# Question: This is commented out because we don't need it? Nodes created in the conditional
# compilation state are captured in the to_case_block? Always?
# Is this still true of nested conditionals? Is that why propeller compiler is complaining?
# branch_nodes = ctx.compilation_state.nodes
node, promises = to_branch_node(self._name, self)
# Verify branch_nodes == nodes in bn
bindings: typing.List[Binding] = []
upstream_nodes = set()
for p in promises:
if not p.is_ready:
bindings.append(
Binding(var=p.var, binding=BindingData(promise=p.ref)))
upstream_nodes.add(p.ref.node)
n = Node(
id=
f"{ctx.compilation_state.prefix}n{len(ctx.compilation_state.nodes)}", # type: ignore
metadata=_core_wf.NodeMetadata(self._name,
timeout=datetime.timedelta(),
retries=RetryStrategy(0)),
bindings=sorted(bindings, key=lambda b: b.var),
upstream_nodes=list(upstream_nodes), # type: ignore
flyte_entity=node,
)
FlyteContextManager.current_context().compilation_state.add_node(
n) # type: ignore
return self._compute_outputs(n)
return self._condition
def test_format_correct():
class TempEncoder(StructuredDatasetEncoder):
def __init__(self):
super().__init__(pd.DataFrame, S3, "avro")
def encode(
self,
ctx: FlyteContext,
structured_dataset: StructuredDataset,
structured_dataset_type: StructuredDatasetType,
) -> literals.StructuredDataset:
return literals.StructuredDataset(
uri="/tmp/avro",
metadata=StructuredDatasetMetadata(structured_dataset_type))
ctx = FlyteContextManager.current_context()
df = pd.DataFrame({"name": ["Tom", "Joseph"], "age": [20, 22]})
annotated_sd_type = Annotated[StructuredDataset, "avro",
kwtypes(name=str, age=int)]
df_literal_type = TypeEngine.to_literal_type(annotated_sd_type)
assert df_literal_type.structured_dataset_type is not None
assert len(df_literal_type.structured_dataset_type.columns) == 2
assert df_literal_type.structured_dataset_type.columns[0].name == "name"
assert df_literal_type.structured_dataset_type.columns[
0].literal_type.simple is not None
assert df_literal_type.structured_dataset_type.columns[1].name == "age"
assert df_literal_type.structured_dataset_type.columns[
1].literal_type.simple is not None
assert df_literal_type.structured_dataset_type.format == "avro"
sd = annotated_sd_type(df)
with pytest.raises(ValueError):
TypeEngine.to_literal(ctx,
sd,
python_type=annotated_sd_type,
expected=df_literal_type)
StructuredDatasetTransformerEngine.register(TempEncoder(),
default_for_type=False)
sd2 = annotated_sd_type(df)
sd_literal = TypeEngine.to_literal(ctx,
sd2,
python_type=annotated_sd_type,
expected=df_literal_type)
assert sd_literal.scalar.structured_dataset.metadata.structured_dataset_type.format == "avro"
@task
def t1() -> Annotated[StructuredDataset, "avro"]:
return StructuredDataset(dataframe=df)
assert t1().file_format == "avro"
def _outputs_interface(self) -> Dict[Any, Variable]:
"""
We override this method from PythonTask because the dispatch_execute method uses this
interface to construct outputs. Each instance of an container_array task will however produce outputs
according to the underlying run_task interface and the array plugin handler will actually create a collection
from these individual outputs as the final output value.
"""
ctx = FlyteContextManager.current_context()
if ctx.execution_state is not None and ctx.execution_state.mode == ExecutionState.Mode.LOCAL_WORKFLOW_EXECUTION:
# In workflow execution mode we actually need to use the parent (mapper) task output interface.
return self.interface.outputs
return self._run_task.interface.outputs
def test_levels():
ctx = FlyteContextManager.current_context()
b = ctx.new_builder()
b.flyte_client = SampleTestClass(value=1)
with FlyteContextManager.with_context(b) as outer:
assert outer.flyte_client.value == 1
b = outer.new_builder()
b.flyte_client = SampleTestClass(value=2)
with FlyteContextManager.with_context(b) as ctx:
assert ctx.flyte_client.value == 2
with FlyteContextManager.with_context(outer.with_new_compilation_state()) as ctx:
assert ctx.flyte_client.value == 1
def __init__(self, lhs: Union["Promise", Any], op: ComparisonOps,
rhs: Union["Promise", Any]):
self._op = op
self._lhs = None
self._rhs = None
if isinstance(lhs, Promise):
self._lhs = lhs
if lhs.is_ready:
if lhs.val.scalar is None or lhs.val.scalar.primitive is None:
raise ValueError(
"Only primitive values can be used in comparison")
if isinstance(rhs, Promise):
self._rhs = rhs
if rhs.is_ready:
if rhs.val.scalar is None or rhs.val.scalar.primitive is None:
raise ValueError(
"Only primitive values can be used in comparison")
if self._lhs is None:
self._lhs = type_engine.TypeEngine.to_literal(
FlyteContextManager.current_context(), lhs, type(lhs), None)
if self._rhs is None:
self._rhs = type_engine.TypeEngine.to_literal(
FlyteContextManager.current_context(), rhs, type(rhs), None)
def current_context() -> ExecutionParameters:
"""
Use this method to get a handle of specific parameters available in a flyte task.
Usage
.. code-block::
flytekit.current_context().logging.info(...)
Available params are documented in :py:class:`flytekit.core.context_manager.ExecutionParams`.
There are some special params, that should be available
"""
return FlyteContextManager.current_context(
).execution_state.user_space_params
def restore(self, path: typing.Optional[typing.Union[Path, str]] = None) -> typing.Optional[Path]:
# We have to lazy load, until we fix the imports
from flytekit.core.context_manager import FlyteContextManager
if self._checkpoint_src is None or self._checkpoint_src == "":
return None
if self._prev_download_path:
return self._prev_download_path
if path is None:
p = Path(self._td.name)
path = p.joinpath(self.SRC_LOCAL_FOLDER)
path.mkdir()
elif isinstance(path, str):
path = Path(path)
if not path.is_dir():
raise ValueError("Checkpoints can be restored to a directory only.")
FlyteContextManager.current_context().file_access.download_directory(self._checkpoint_src, str(path))
self._prev_download_path = path
return self._prev_download_path
def start_branch(self, c: Case, last_case: bool = False) -> Case:
"""
At the start of an execution of every branch this method should be called.
:param c: -> the case that represents this branch
:param last_case: -> a boolean that indicates if this is the last branch in the ifelseblock
"""
added_case = super().start_branch(c, last_case)
ctx = FlyteContextManager.current_context()
# This is a short-circuit for the case when the branch was taken
# We already have a candidate case selected
if self._selected_case is None:
if c.expr is None or c.expr.eval() or last_case:
ctx.execution_state.take_branch()
self._selected_case = added_case
return added_case
def __call__(self, *args, **kwargs):
# When a Task is () aka __called__, there are three things we may do:
# a. Task Execution Mode - just run the Python function as Python normally would. Flyte steps completely
# out of the way.
# b. Compilation Mode - this happens when the function is called as part of a workflow (potentially
# dynamic task?). Instead of running the user function, produce promise objects and create a node.
# c. Workflow Execution Mode - when a workflow is being run locally. Even though workflows are functions
# and everything should be able to be passed through naturally, we'll want to wrap output values of the
# function into objects, so that potential .with_cpu or other ancillary functions can be attached to do
# nothing. Subsequent tasks will have to know how to unwrap these. If by chance a non-Flyte task uses a
# task output as an input, things probably will fail pretty obviously.
if len(args) > 0:
raise _user_exceptions.FlyteAssertion(
f"When calling tasks, only keyword args are supported. "
f"Aborting execution as detected {len(args)} positional args {args}"
)
ctx = FlyteContextManager.current_context()
if ctx.compilation_state is not None and ctx.compilation_state.mode == 1:
return self.compile(ctx, *args, **kwargs)
elif (ctx.execution_state is not None and ctx.execution_state.mode
== ExecutionState.Mode.LOCAL_WORKFLOW_EXECUTION):
if ctx.execution_state.branch_eval_mode == BranchEvalMode.BRANCH_SKIPPED:
if self.python_interface and self.python_interface.output_tuple_name:
variables = [
k for k in self.python_interface.outputs.keys()
]
output_tuple = collections.namedtuple(
self.python_interface.output_tuple_name, variables)
nones = [
None for _ in self.python_interface.outputs.keys()
]
return output_tuple(*nones)
else:
# Should we return multiple None's here?
return None
return self._local_execute(ctx, **kwargs)
else:
logger.warning("task run without context - executing raw function")
new_user_params = self.pre_execute(ctx.user_space_params)
with FlyteContextManager.with_context(
ctx.with_execution_state(
ctx.execution_state.with_params(
mode=ExecutionState.Mode.LOCAL_TASK_EXECUTION,
user_space_params=new_user_params))):
return self.execute(**kwargs)