def get_serializable_task(
entity_mapping: OrderedDict,
settings: SerializationSettings,
entity: FlyteLocalEntity,
) -> task_models.TaskSpec:
task_id = _identifier_model.Identifier(
_identifier_model.ResourceType.TASK,
settings.project,
settings.domain,
entity.name,
settings.version,
)
if settings.should_fast_serialize() and isinstance(entity, PythonAutoContainerTask):
# For fast registration, we'll need to muck with the command, but only for certain kinds of tasks. Specifically,
# tasks that rely on user code defined in the container. This should be encapsulated by the auto container
# parent class
entity.set_command_fn(_fast_serialize_command_fn(settings, entity))
tt = task_models.TaskTemplate(
id=task_id,
type=entity.task_type,
metadata=entity.metadata.to_taskmetadata_model(),
interface=entity.interface,
custom=entity.get_custom(settings),
container=entity.get_container(settings),
task_type_version=entity.task_type_version,
security_context=entity.security_context,
config=entity.get_config(settings),
k8s_pod=entity.get_k8s_pod(settings),
)
if settings.should_fast_serialize() and isinstance(entity, PythonAutoContainerTask):
entity.reset_command_fn()
return task_models.TaskSpec(template=tt)
def register(self, project, domain, name, version):
"""
:param Text project: The project in which to register this task.
:param Text domain: The domain in which to register this task.
:param Text name: The name to give this task.
:param Text version: The version in which to register this task.
"""
# TODO: Revisit the notion of supplying the project, domain, name, version, as opposed to relying on the
# current ID.
self.validate()
id_to_register = _identifier.Identifier(
_identifier_model.ResourceType.TASK, project, domain, name,
version)
old_id = self.id
client = _flyte_engine.get_client()
try:
self._id = id_to_register
client.create_task(id_to_register, _task_model.TaskSpec(self))
self._id = old_id
self._has_registered = True
return str(id_to_register)
except _user_exceptions.FlyteEntityAlreadyExistsException:
pass
except Exception:
self._id = old_id
raise
def register(self, identifier):
client = _FlyteClientManager(
_platform_config.URL.get(),
insecure=_platform_config.INSECURE.get()).client
try:
client.create_task(identifier,
_task_models.TaskSpec(self.sdk_task))
except _user_exceptions.FlyteEntityAlreadyExistsException:
pass
def get_serializable_task(
entity_mapping: OrderedDict,
settings: SerializationSettings,
entity: FlyteLocalEntity,
fast: bool,
) -> task_models.TaskSpec:
task_id = _identifier_model.Identifier(
_identifier_model.ResourceType.TASK,
settings.project,
settings.domain,
entity.name,
settings.version,
)
tt = task_models.TaskTemplate(
id=task_id,
type=entity.task_type,
metadata=entity.metadata.to_taskmetadata_model(),
interface=entity.interface,
custom=entity.get_custom(settings),
container=entity.get_container(settings),
task_type_version=entity.task_type_version,
security_context=entity.security_context,
config=entity.get_config(settings),
k8s_pod=entity.get_k8s_pod(settings),
)
# For fast registration, we'll need to muck with the command, but only for certain kinds of tasks. Specifically,
# tasks that rely on user code defined in the container. This should be encapsulated by the auto container
# parent class
if fast and isinstance(entity, PythonAutoContainerTask):
args = [
"pyflyte-fast-execute",
"--additional-distribution",
"{{ .remote_package_path }}",
"--dest-dir",
"{{ .dest_dir }}",
"--",
] + tt.container.args[:]
del tt.container.args[:]
tt.container.args.extend(args)
return task_models.TaskSpec(template=tt)
def serialize(self):
"""
:rtype: flyteidl.admin.task_pb2.TaskSpec
"""
return _task_model.TaskSpec(self).to_flyte_idl()
def get_serializable_task(
entity_mapping: OrderedDict,
settings: SerializationSettings,
entity: FlyteLocalEntity,
) -> task_models.TaskSpec:
task_id = _identifier_model.Identifier(
_identifier_model.ResourceType.TASK,
settings.project,
settings.domain,
entity.name,
settings.version,
)
if isinstance(
entity, PythonFunctionTask
) and entity.execution_mode == PythonFunctionTask.ExecutionBehavior.DYNAMIC:
# In case of Dynamic tasks, we want to pass the serialization context, so that they can reconstruct the state
# from the serialization context. This is passed through an environment variable, that is read from
# during dynamic serialization
settings = settings.with_serialized_context()
container = entity.get_container(settings)
# This pod will be incorrect when doing fast serialize
pod = entity.get_k8s_pod(settings)
if settings.should_fast_serialize():
# This handles container tasks.
if container and isinstance(entity,
(PythonAutoContainerTask, MapPythonTask)):
# For fast registration, we'll need to muck with the command, but on
# ly for certain kinds of tasks. Specifically,
# tasks that rely on user code defined in the container. This should be encapsulated by the auto container
# parent class
container._args = prefix_with_fast_execute(settings,
container.args)
# If the pod spec is not None, we have to get it again, because the one we retrieved above will be incorrect.
# The reason we have to call get_k8s_pod again, instead of just modifying the command in this file, is because
# the pod spec is a K8s library object, and we shouldn't be messing around with it in this file.
elif pod:
if isinstance(entity, MapPythonTask):
entity.set_command_prefix(
get_command_prefix_for_fast_execute(settings))
pod = entity.get_k8s_pod(settings)
else:
entity.set_command_fn(
_fast_serialize_command_fn(settings, entity))
pod = entity.get_k8s_pod(settings)
entity.reset_command_fn()
tt = task_models.TaskTemplate(
id=task_id,
type=entity.task_type,
metadata=entity.metadata.to_taskmetadata_model(),
interface=entity.interface,
custom=entity.get_custom(settings),
container=container,
task_type_version=entity.task_type_version,
security_context=entity.security_context,
config=entity.get_config(settings),
k8s_pod=pod,
sql=entity.get_sql(settings),
)
if settings.should_fast_serialize() and isinstance(
entity, PythonAutoContainerTask):
entity.reset_command_fn()
return task_models.TaskSpec(template=tt)