Repository: piskvorky/smart_open

Issue:
use prefetch for sftp?
Hi,

I'm finding certain files slow to read over sftp. debugging lead me to wonder if we should be calling `prefetch()` before returning the file handle

https://docs.paramiko.org/en/latest/api/sftp.html#paramiko.sftp_file.SFTPFile.prefetch

Thoughts?


Discussion:
> Pre-fetch the remaining contents of this file in anticipation of future [read](https://docs.paramiko.org/en/latest/api/sftp.html#paramiko.sftp_file.SFTPFile.read) calls.

that sounds to me like the *entire* file will be downloaded when calling prefetch, which would be an anti-pattern since this lib is designed for files that do not fit into memory
while i can't find it in the docs, from testing with large (~1gb) files paramiko doesn't appear to load the entire file into memory, only a fixed buffer

i was trying to confirm this in the paramiko code but couldn't quite follow it
fwiw, `smart_open.s3` also uses a prefetch mechanic using `self._buffer.fill`: https://github.com/piskvorky/smart_open/blob/76415ca7b64f64a4be6a926433d03243936d2811/smart_open/s3.py#L687

I'm not a maintainer (just an active contributor), but if the API is analogous to that implementation (global `DEFAULT_BUFFER_SIZE` variable etc) so the memory usage can be controlled transparently, I think it would have a good chance of landing :)

cc @mpenkov 
If this does not change smart_open's API significantly, I'm open to a PR.
ok i managed to understand paramiko's code a bit better and i'm now pretty sure only at most a single chunk (32k) is stored in memory

regarding api, i expect we can handle it with `transport_params`. indeed, i think the docs https://github.com/piskvorky/smart_open/blob/develop/help.txt#L187C12-L187C79 (`Any additional settings to be passed to paramiko.SSHClient.connect`) are fortunately incorrect, only `transport_params['connect_kwargs']` are passed to `connect` which leaves space to add say `prefetch_kwargs` or similar (name suggestions welcome).

as a first suggestion, how about the presence of `prefetch_kwargs` (dict) will trigger prefetching (to eg keep backwargs compat) and the dict itself is passed as kwargs to the `prefetch` call (to eg control max concurrency)

the slight downside is that it makes this (pretty significant) performance increase opt-in, but on the other hand it's safer to add. making prefetch the default would presumably require some other transport param value to disable it again

if this suggestion sounds ok i'm happy to make a pr
if we're talking about 32k prefetch, probably there's not much of a downside to having this enabled permanently? can you permalink the relevant paramiko code for the prefetch size?
at most one chunk stays in memory https://github.com/paramiko/paramiko/blob/ed8b09751ff20340332d4b1bb2b10e32aedc57ff/paramiko/sftp_file.py#L149
(added if missing, deleted when read (put back if partially read)

chunks are made here https://github.com/paramiko/paramiko/blob/ed8b09751ff20340332d4b1bb2b10e32aedc57ff/paramiko/sftp_file.py#L478
max size definition https://github.com/paramiko/paramiko/blob/ed8b09751ff20340332d4b1bb2b10e32aedc57ff/paramiko/sftp_file.py#L59
though while cheap on the client, it does use a bunch of resources on the server which can be wasteful if the intention is to only read a small part of the file. so maybe we do want it opt-in
ok i'm not sure why i didn't realise this sooner but if i need custom code for a given protocol (to specify custom transport params) i guess i could also just call `prefetch()` on the returned file object. or am i missing something?
you'll be talking to this object:
https://github.com/piskvorky/smart_open/blob/76415ca7b64f64a4be6a926433d03243936d2811/smart_open/smart_open_lib.py#L263
which proxies to your paramiko object only if there are no other proxies in front of it like a decompression handler
ok. raised #849 
to me it still looks like paramiko fetches the entire file into memory when you call `prefetch`: 
https://github.com/paramiko/paramiko/blob/ed8b09751ff20340332d4b1bb2b10e32aedc57ff/paramiko/sftp_file.py#L482

it's not prefetching lazy/just-in-time like with `smart_open.s3`. could that be achieved with a more low-level approach?
my understanding is that it _requests all the data from the server (hence my comment about server resources above) but it only reads the replies as required and only keeps a single such chunk of data in memory at a time
aha, that's not obvious for me from a quick look at the code. would be cool to add a test that verifies this behaviour: when a big file is opened with `prefetch_kwargs` passed, some **but not all** data gets immediately requested from the server even before performing any read operation
>  some but not all data gets immediately requested from the serve

all the data does get requested but none of it read until we start reading

are there any existing tests that make real connections like this that i can start from?
I don't know if I'm overthinking this, and I don't know if such test is even necessary. probably for @mpenkov to decide if he's comfortable with this :)
Hey guys, I'm using very smart_open library to read large files from SFTP server and I observed that the reads are very slow. It is taking 2 sec to read 8 MB chunks form the file. I think adding support for prefetch would help. Please see https://docs.paramiko.org/en/latest/api/sftp.html 

I'm waiting on this MR to get merged in. If it is going to take more time to get this merged, I will try switching over to use paramiko. Please let me know if any one has updates on when this can be merged.
@nithinkrishna2107 it would be great if you can test the PR and report back!

in pip / setup.py / pyproject.toml:
```
smart_open[ssh]@https://github.com/davidszotten/smart_open/archive/refs/heads/sftp-prefetch.zip
```
or (equivalent but immutable):
```
smart_open[ssh]@https://github.com/piskvorky/smart_open/archive/afe7da9fb0632bc4f6686572c22f1fc83bc9b550.zip
```
> in pip / setup.py / pyproject.toml:
> 
> ```
> smart_open[ssh]@https://github.com/davidszotten/smart_open/archive/refs/heads/sftp-prefetch.zip
> ```

Hi, FYI I've tested in a staging environment using the PR code, as you suggested, for an sftp compressed file download that would previously take over **35 minutes** with smart_open to stream from the sftp server onto the local filesystem.

The only change (aside from installing the PR version of smart_open) is:

```python
transport_params = {"prefetch_kwargs": {}}
with smart_open.open(url, "rb", transport_params=transport_params) as file_handle:
    # rest of the code
```

And the file now downloads in less than **30 seconds** which is much better. 
@ddelange Sure, I will test this and let you know the results. 

Thanks to @tvs-lyst for testing the change. 

I observed that smart_open has 2 files where open is defined one here https://github.com/piskvorky/smart_open/blob/develop/smart_open/smart_open_lib.py#L100 and one here https://github.com/piskvorky/smart_open/blob/develop/smart_open/ssh.py

Do we need to make similar change in smart_open_lib.py?
no need to touch smart_open_lib (from smart_open import open), that one just propagates transport_params to (in this case) smart_open.ssh.open. so #849 is sufficient
@ddelange Thanks so much for the fix. I have tested the fix with 2GB file and it took 1 min to read the file which is 8x improvement in read speed.
thanks to @davidszotten :)
yeah thanks to @davidszotten :-)
@davidszotten @ddelange Any idea when this can be merged? 



Patch:
diff --git a/smart_open/ssh.py b/smart_open/ssh.py
index 1dc73cbe..3af25064 100644
--- a/smart_open/ssh.py
+++ b/smart_open/ssh.py
@@ -111,21 +111,18 @@ def parse_uri(uri_as_string):
 
 
 def open_uri(uri, mode, transport_params):
-    # `connect_kwargs` is a legitimate param *only* for sftp, so this filters it out of validation
-    #   (otherwise every call with this present complains it's not valid)
-    params_to_validate = {k: v for k, v in transport_params.items() if k != 'connect_kwargs'}
-    smart_open.utils.check_kwargs(open, params_to_validate)
+    kwargs = smart_open.utils.check_kwargs(open, transport_params)
     parsed_uri = parse_uri(uri)
     uri_path = parsed_uri.pop('uri_path')
     parsed_uri.pop('scheme')
-    return open(uri_path, mode, transport_params=transport_params, **parsed_uri)
+    return open(uri_path, mode, **parsed_uri, **kwargs)
 
 
-def _connect_ssh(hostname, username, port, password, transport_params):
+def _connect_ssh(hostname, username, port, password, connect_kwargs):
     ssh = paramiko.SSHClient()
     ssh.load_system_host_keys()
     ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
-    kwargs = transport_params.get('connect_kwargs', {}).copy()
+    kwargs = (connect_kwargs or {}).copy()
     if 'key_filename' not in kwargs:
         kwargs.setdefault('password', password)
     kwargs.setdefault('username', username)
@@ -133,17 +130,13 @@ def _connect_ssh(hostname, username, port, password, transport_params):
     return ssh
 
 
-def _maybe_fetch_config(host, username=None, password=None, port=None, transport_params=None):
+def _maybe_fetch_config(host, username=None, password=None, port=None, connect_kwargs=None):
     # If all fields are set, return as-is.
-    if not any(arg is None for arg in (host, username, password, port, transport_params)):
-        return host, username, password, port, transport_params
+    if not any(arg is None for arg in (host, username, password, port, connect_kwargs)):
+        return host, username, password, port, connect_kwargs
 
     if not host:
         raise ValueError('you must specify the host to connect to')
-    if not transport_params:
-        transport_params = {}
-    if "connect_kwargs" not in transport_params:
-        transport_params["connect_kwargs"] = {}
 
     # Attempt to load an OpenSSH config.
     #
@@ -160,7 +153,7 @@ def _maybe_fetch_config(host, username=None, password=None, port=None, transport
     # - compression selection
     # - GSS configuration
     #
-    connect_params = transport_params["connect_kwargs"]
+    connect_params = (connect_kwargs or {}).copy()
     config_files = [f for f in _SSH_CONFIG_FILES if os.path.exists(f)]
     #
     # This is the actual name of the host.  The input host may actually be an
@@ -221,10 +214,19 @@ def _maybe_fetch_config(host, username=None, password=None, port=None, transport
     if actual_hostname:
         host = actual_hostname
 
-    return host, username, password, port, transport_params
+    return host, username, password, port, connect_params
 
 
-def open(path, mode='r', host=None, user=None, password=None, port=None, transport_params=None):
+def open(
+    path,
+    mode="r",
+    host=None,
+    user=None,
+    password=None,
+    port=None,
+    connect_kwargs=None,
+    prefetch_kwargs=None,
+):
     """Open a file on a remote machine over SSH.
 
     Expects authentication to be already set up via existing keys on the local machine.
@@ -244,8 +246,11 @@ def open(path, mode='r', host=None, user=None, password=None, port=None, transpo
         The password to use to login to the remote machine.
     port: int, optional
         The port to connect to.
-    transport_params: dict, optional
-        Any additional settings to be passed to paramiko.SSHClient.connect
+    connect_kwargs: dict, optional
+        Any additional settings to be passed to paramiko.SSHClient.connect.
+    prefetch_kwargs: dict, optional
+        Any additional settings to be passed to paramiko.SFTPFile.prefetch.
+        The presence of this dict (even if empty) triggers prefetching.
 
     Returns
     -------
@@ -259,8 +264,8 @@ def open(path, mode='r', host=None, user=None, password=None, port=None, transpo
     If ``username`` or ``password`` are specified in *both* the uri and
     ``transport_params``, ``transport_params`` will take precedence
     """
-    host, user, password, port, transport_params = _maybe_fetch_config(
-        host, user, password, port, transport_params
+    host, user, password, port, connect_kwargs = _maybe_fetch_config(
+        host, user, password, port, connect_kwargs
     )
 
     key = (host, user)
@@ -273,9 +278,9 @@ def open(path, mode='r', host=None, user=None, password=None, port=None, transpo
             #   and if not, refresh the connection
             if not ssh.get_transport().active:
                 ssh.close()
-                ssh = _SSH[key] = _connect_ssh(host, user, port, password, transport_params)
+                ssh = _SSH[key] = _connect_ssh(host, user, port, password, connect_kwargs)
         except KeyError:
-            ssh = _SSH[key] = _connect_ssh(host, user, port, password, transport_params)
+            ssh = _SSH[key] = _connect_ssh(host, user, port, password, connect_kwargs)
 
         try:
             transport = ssh.get_transport()
@@ -294,4 +299,6 @@ def open(path, mode='r', host=None, user=None, password=None, port=None, transpo
 
     fobj = sftp_client.open(path, mode)
     fobj.name = path
+    if prefetch_kwargs is not None:
+        fobj.prefetch(**prefetch_kwargs)
     return fobj


Test patch:
diff --git a/smart_open/tests/test_smart_open.py b/smart_open/tests/test_smart_open.py
index e37c1aee..45fd7d86 100644
--- a/smart_open/tests/test_smart_open.py
+++ b/smart_open/tests/test_smart_open.py
@@ -474,7 +474,6 @@ def test_read_ssh(self, mock_open):
             user='ubuntu',
             password='pass',
             port=1022,
-            transport_params={'hello': 'world'},
         )
 
     @responses.activate
diff --git a/smart_open/tests/test_ssh.py b/smart_open/tests/test_ssh.py
index c6c291a0..7b99d98f 100644
--- a/smart_open/tests/test_ssh.py
+++ b/smart_open/tests/test_ssh.py
@@ -127,6 +127,23 @@ def test_open_with_openssh_config_override_user(self, mock_connect, get_transp_m
             gss_trust_dns=False,
         )
 
+    @mock_ssh
+    def test_open_with_prefetch(self, mock_connect, get_transp_mock):
+        smart_open.open(
+            "ssh://user:pass@some-host/",
+            transport_params={"prefetch_kwargs": {"max_concurrent_requests": 3}},
+        )
+        mock_sftp = get_transp_mock().open_sftp_client()
+        mock_fobj = mock_sftp.open()
+        mock_fobj.prefetch.assert_called_with(max_concurrent_requests=3)
+
+    @mock_ssh
+    def test_open_without_prefetch(self, mock_connect, get_transp_mock):
+        smart_open.open("ssh://user:pass@some-host/")
+        mock_sftp = get_transp_mock().open_sftp_client()
+        mock_fobj = mock_sftp.open()
+        mock_fobj.prefetch.assert_not_called()
+
 
 if __name__ == "__main__":
     logging.basicConfig(format="%(asctime)s : %(levelname)s : %(message)s", level=logging.DEBUG)