[Artifacts/Elastic] Skip partitions

[wild-endeavor]

Why are the changes needed?

This relates to an issue we found when using the artifact class with Elastic while also specifying partitions.

As far as I can tell, what's happening is that the kfpytorch plugin serializes, using cloud pickle, and then deserializes within the worker. Because of how cloudpickle works, this was causing an infinite recursion in the __getattr__ call of the partitions object. Basically the getattr function is called only if the field is missing from the object, but cloudpickle doesn't run the init function for the object, which leads to an infinite recursion. The added unit test fails with infinite recursion if the code change to getattr is removed. (See more background here.)

  File "/Users/ytong/go/src/github.com/flyteorg/flytekit/plugins/flytekit-kf-pytorch/flytekitplugins/kfpytorch/task.py", line 276, in spawn_helper
    fn = cloudpickle.loads(fn)
         ^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ytong/go/src/github.com/flyteorg/flytekit/flytekit/core/artifact.py", line 321, in __getattr__
    if self.partitions and item in self.partitions:
       ^^^^^^^^^^^^^^^
  File "/Users/ytong/go/src/github.com/flyteorg/flytekit/flytekit/core/artifact.py", line 312, in partitions
    return self._partitions
           ^^^^^^^^^^^^^^^^
  File "/Users/ytong/go/src/github.com/flyteorg/flytekit/flytekit/core/artifact.py", line 321, in __getattr__
    if self.partitions and item in self.partitions:
       ^^^^^^^^^^^^^^^

This was hard to find because the exception was getting swallowed by Elastic somewhere. This stack trace was taken from a file which looks something like /var/folders/jx/54tww2ls58n8qtlp9k31nbd80000gp/T/torchelastic_wg0jjz6j/local_r5pr8f9n/attempt_0/0/error.json. The behavior of the pytorch plugin is such that this just hangs (at least until some timeout), rather than failing even though there is an error.json.

What changes were proposed in this pull request?

How was this patch tested?

Tested locally.

Setup process

Screenshots

Check all the applicable boxes

• I updated the documentation accordingly.
• All new and existing tests passed.
• All commits are signed-off.

Related PRs

Docs link

[samhita-alla]

@fg91 do you know why the plugin doesn't fail immediately when there's an error.json file?

This was hard to find because the exception was getting swallowed by Elastic somewhere. This stack trace was taken from a file which looks something like /var/folders/jx/54tww2ls58n8qtlp9k31nbd80000gp/T/torchelastic_wg0jjz6j/local_r5pr8f9n/attempt_0/0/error.json. The behavior of the pytorch plugin is such that this just hangs (at least until some timeout), rather than failing even though there is an error.json.

[fg91]

Nit: Could we please add om to ElasticWorkerResult's doc string? 🙇

[fg91]

@fg91 do you know why the plugin doesn't fail immediately when there's an error.json file?

I don't know unfortunately, sorry. But I can confirm that this is not the normal behavior, normally it fails immediately:

The behavior of the pytorch plugin is such that this just hangs (at least until some timeout), rather than failing even though there is an error.json.