Skip to content

Commit

Permalink
Improve automatic file transfer strategy selection
Browse files Browse the repository at this point in the history
Handle the following cases:
- Local-Remote (except HTTP destination)
- Remote-Local
- Local-Local

This will be needed for to port remote executors to the
new runner-executor API.
  • Loading branch information
cjao committed May 16, 2024
1 parent 5486a14 commit f2f0632
Show file tree
Hide file tree
Showing 3 changed files with 48 additions and 8 deletions.
3 changes: 3 additions & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -92,3 +92,6 @@ node_modules/

# Ignore mock database
**/*.sqlite

# Ignore virtual envs
*.venv
36 changes: 28 additions & 8 deletions covalent/_file_transfer/file_transfer.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,10 +18,37 @@

from .enums import FileTransferStrategyTypes, FtCallDepReturnValue, Order
from .file import File
from .strategies.gcloud_strategy import GCloud
from .strategies.http_strategy import HTTP
from .strategies.s3_strategy import S3
from .strategies.shutil_strategy import Shutil
from .strategies.transfer_strategy_base import FileTransferStrategy

# TODO: make this pluggable similar to executor plugins
_strategy_type_map = {
FileTransferStrategyTypes.Shutil: Shutil,
FileTransferStrategyTypes.S3: S3,
FileTransferStrategyTypes.HTTP: HTTP,
FileTransferStrategyTypes.GCloud: GCloud,
}


def _guess_transfer_strategy(from_file: File, to_file: File) -> FileTransferStrategy:
# Handle the following cases automatically
# Local-Remote (except HTTP destination)
# Remote-local
# Local-local

if (
from_file.mapped_strategy_type == FileTransferStrategyTypes.Shutil
and to_file.mapped_strategy_type != FileTransferStrategyTypes.HTTP
):
return _strategy_type_map[to_file.mapped_strategy_type]
elif to_file.mapped_strategy_type == FileTransferStrategyTypes.Shutil:
return _strategy_type_map[from_file.mapped_strategy_type]
else:
raise AttributeError("FileTransfer requires a file transfer strategy to be specified")


class FileTransfer:
"""
Expand Down Expand Up @@ -58,15 +85,8 @@ def __init__(
# assign explicit strategy or default to strategy based on from_file & to_file schemes
if strategy:
self.strategy = strategy
elif (
from_file.mapped_strategy_type == FileTransferStrategyTypes.Shutil
and to_file.mapped_strategy_type == FileTransferStrategyTypes.Shutil
):
self.strategy = Shutil()
elif from_file.mapped_strategy_type == FileTransferStrategyTypes.HTTP:
self.strategy = HTTP()
else:
raise AttributeError("FileTransfer requires a file transfer strategy to be specified")
self.strategy = _guess_transfer_strategy(from_file, to_file)()

self.to_file = to_file
self.from_file = from_file
Expand Down
17 changes: 17 additions & 0 deletions tests/covalent_tests/file_transfer/file_transfer_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,8 @@
TransferToRemote,
)
from covalent._file_transfer.strategies.rsync_strategy import Rsync
from covalent._file_transfer.strategies.s3_strategy import S3
from covalent._file_transfer.strategies.shutil_strategy import Shutil


class TestFileTransfer:
Expand Down Expand Up @@ -109,3 +111,18 @@ def test_transfer_to_remote(self):

with pytest.raises(ValueError):
result = TransferToRemote("file:///home/one", "file:///home/one/", strategy=strategy)

def test_auto_transfer_strategy(self):
from_file = File("s3://bucket/object.pkl")
to_file = File("file:///tmp/object.pkl")
ft = FileTransfer(from_file, to_file)
assert type(ft.strategy) is S3

ft = FileTransfer(to_file, from_file)
assert type(ft.strategy) is S3

ft = FileTransfer(to_file, to_file)
assert type(ft.strategy) is Shutil

with pytest.raises(AttributeError):
_ = FileTransfer(from_file, from_file)

0 comments on commit f2f0632

Please sign in to comment.