Skip to content

Commit

Permalink
Fix Base Agent Bug (#2057)
Browse files Browse the repository at this point in the history
Signed-off-by: Future Outlier <eric901201@gmai.com>
Co-authored-by: Future Outlier <eric901201@gmai.com>
  • Loading branch information
Future-Outlier and Future Outlier authored Dec 16, 2023
1 parent 7f7272f commit 628ac89
Showing 1 changed file with 6 additions and 3 deletions.
9 changes: 6 additions & 3 deletions flytekit/extend/backend/base_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -204,17 +204,20 @@ async def _create(
literals = inputs or {}
for k, v in inputs.items():
literals[k] = TypeEngine.to_literal(ctx, v, type(v), self._entity.interface.inputs[k].type)

literal_map = LiteralMap(literals)

if isinstance(self, PythonFunctionTask):
# Write the inputs to a remote file, so that the remote task can read the inputs from this file.
path = ctx.file_access.get_random_local_path()
utils.write_proto_to_file(LiteralMap(literals).to_flyte_idl(), path)
utils.write_proto_to_file(literal_map.to_flyte_idl(), path)
ctx.file_access.put_data(path, f"{output_prefix}/inputs.pb")
task_template = render_task_template(task_template, output_prefix)

if self._agent.asynchronous:
res = await self._agent.async_create(grpc_ctx, output_prefix, task_template, inputs)
res = await self._agent.async_create(grpc_ctx, output_prefix, task_template, literal_map)
else:
res = self._agent.create(grpc_ctx, output_prefix, task_template, inputs)
res = self._agent.create(grpc_ctx, output_prefix, task_template, literal_map)

signal.signal(signal.SIGINT, partial(self.signal_handler, res.resource_meta)) # type: ignore
return res
Expand Down

0 comments on commit 628ac89

Please sign in to comment.