Skip to content

Commit

Permalink
Merge pull request #522 from swarajpande5/relocate-fix
Browse files Browse the repository at this point in the history
Relocate `File` and `HttpClient` to their respective modules
  • Loading branch information
ericcurtin authored Dec 22, 2024
2 parents 3c8c0bd + f2e4d1c commit f3312c9
Show file tree
Hide file tree
Showing 3 changed files with 188 additions and 177 deletions.
178 changes: 1 addition & 177 deletions ramalama/common.py
Original file line number Diff line number Diff line change
@@ -1,16 +1,15 @@
"""ramalama common module."""

import fcntl
import hashlib
import os
import random
import shutil
import string
import subprocess
import sys
import time
import urllib.request
import urllib.error
from ramalama.http_client import HttpClient

x = False
mnt_dir = "/mnt/models"
Expand Down Expand Up @@ -155,181 +154,6 @@ def genname():
return "ramalama_" + "".join(random.choices(string.ascii_letters + string.digits, k=10))


# The following code is inspired from: https://github.com/ericcurtin/lm-pull/blob/main/lm-pull.py

class File:
def __init__(self):
self.file = None
self.fd = -1

def open(self, filename, mode):
self.file = open(filename, mode)
return self.file

def lock(self):
if self.file:
self.fd = self.file.fileno()
try:
fcntl.flock(self.fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
except BlockingIOError:
self.fd = -1
return 1

return 0

def __del__(self):
if self.fd >= 0:
fcntl.flock(self.fd, fcntl.LOCK_UN)

if self.file:
self.file.close()


class HttpClient:
def __init__(self):
pass

def init(self, url, headers, output_file, progress, response_str=None):
output_file_partial = None
if output_file:
output_file_partial = output_file + ".partial"

self.file_size = self.set_resume_point(output_file_partial)
self.printed = False
if self.urlopen(url, headers):
return 1

self.total_to_download = int(self.response.getheader('content-length', 0))
if response_str is not None:
response_str.append(self.response.read().decode('utf-8'))
else:
out = File()
if not out.open(output_file_partial, "ab"):
print("Failed to open file")

return 1

if out.lock():
print("Failed to exclusively lock file")

return 1

self.now_downloaded = 0
self.start_time = time.time()
self.perform_download(out.file, progress)

if output_file:
os.rename(output_file_partial, output_file)

if self.printed:
print("\n")

return 0

def urlopen(self, url, headers):
headers["Range"] = f"bytes={self.file_size}-"
request = urllib.request.Request(url, headers=headers)
try:
self.response = urllib.request.urlopen(request)
except urllib.error.HTTPError as e:
print(f"Request failed: {e.code}", file=sys.stderr)

return 1

if self.response.status not in (200, 206):
print(f"Request failed: {self.response.status}", file=sys.stderr)

return 1

return 0

def perform_download(self, file, progress):
self.total_to_download += self.file_size
self.now_downloaded = 0
self.start_time = time.time()
while True:
data = self.response.read(1024)
if not data:
break

size = file.write(data)
if progress:
self.update_progress(size)

def human_readable_time(self, seconds):
hrs = int(seconds) // 3600
mins = (int(seconds) % 3600) // 60
secs = int(seconds) % 60
width = 10
if hrs > 0:
return f"{hrs}h {mins:02}m {secs:02}s".rjust(width)
elif mins > 0:
return f"{mins}m {secs:02}s".rjust(width)
else:
return f"{secs}s".rjust(width)

def human_readable_size(self, size):
width = 10
for unit in ["B", "KB", "MB", "GB", "TB"]:
if size < 1024:
return f"{size:.2f} {unit}".rjust(width)

size /= 1024

return f"{size:.2f} PB".rjust(width)

def get_terminal_width(self):
return shutil.get_terminal_size().columns

def generate_progress_prefix(self, percentage):
return f"{percentage}% |".rjust(6)

def generate_progress_suffix(self, now_downloaded_plus_file_size, speed, estimated_time):
return f"{self.human_readable_size(now_downloaded_plus_file_size)}/{self.human_readable_size(self.total_to_download)}{self.human_readable_size(speed)}/s{self.human_readable_time(estimated_time)}" # noqa: E501

def calculate_progress_bar_width(self, progress_prefix, progress_suffix):
progress_bar_width = self.get_terminal_width() - len(progress_prefix) - len(progress_suffix) - 3
if progress_bar_width < 1:
progress_bar_width = 1

return progress_bar_width

def generate_progress_bar(self, progress_bar_width, percentage):
pos = (percentage * progress_bar_width) // 100
progress_bar = ""
for i in range(progress_bar_width):
progress_bar += "█" if i < pos else " "

return progress_bar

def set_resume_point(self, output_file):
if output_file and os.path.exists(output_file):
return os.path.getsize(output_file)

return 0

def print_progress(self, progress_prefix, progress_bar, progress_suffix):
print(f"\r{progress_prefix}{progress_bar}| {progress_suffix}", end="")

def update_progress(self, chunk_size):
self.now_downloaded += chunk_size
now_downloaded_plus_file_size = self.now_downloaded + self.file_size
percentage = (now_downloaded_plus_file_size * 100) // self.total_to_download if self.total_to_download else 100
progress_prefix = self.generate_progress_prefix(percentage)
speed = self.calculate_speed(self.now_downloaded, self.start_time)
tim = (self.total_to_download - self.now_downloaded) // speed
progress_suffix = self.generate_progress_suffix(now_downloaded_plus_file_size, speed, tim)
progress_bar_width = self.calculate_progress_bar_width(progress_prefix, progress_suffix)
progress_bar = self.generate_progress_bar(progress_bar_width, percentage)
self.print_progress(progress_prefix, progress_bar, progress_suffix)
self.printed = True

def calculate_speed(self, now_downloaded, start_time):
now = time.time()
elapsed_seconds = now - start_time
return now_downloaded / elapsed_seconds


def download_file(url, dest_path, headers=None, show_progress=True):
"""
Downloads a file from a given URL to a specified destination path.
Expand Down
31 changes: 31 additions & 0 deletions ramalama/file.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
# The following code is inspired from: https://github.com/ericcurtin/lm-pull/blob/main/lm-pull.py

import fcntl


class File:
def __init__(self):
self.file = None
self.fd = -1

def open(self, filename, mode):
self.file = open(filename, mode)
return self.file

def lock(self):
if self.file:
self.fd = self.file.fileno()
try:
fcntl.flock(self.fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
except BlockingIOError:
self.fd = -1
return 1

return 0

def __del__(self):
if self.fd >= 0:
fcntl.flock(self.fd, fcntl.LOCK_UN)

if self.file:
self.file.close()
156 changes: 156 additions & 0 deletions ramalama/http_client.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,156 @@
# The following code is inspired from: https://github.com/ericcurtin/lm-pull/blob/main/lm-pull.py

import os
import shutil
import sys
import time
import urllib.request
import urllib.error
from ramalama.file import File


class HttpClient:
def __init__(self):
pass

def init(self, url, headers, output_file, progress, response_str=None):
output_file_partial = None
if output_file:
output_file_partial = output_file + ".partial"

self.file_size = self.set_resume_point(output_file_partial)
self.printed = False
if self.urlopen(url, headers):
return 1

self.total_to_download = int(self.response.getheader('content-length', 0))
if response_str is not None:
response_str.append(self.response.read().decode('utf-8'))
else:
out = File()
if not out.open(output_file_partial, "ab"):
print("Failed to open file")

return 1

if out.lock():
print("Failed to exclusively lock file")

return 1

self.now_downloaded = 0
self.start_time = time.time()
self.perform_download(out.file, progress)

if output_file:
os.rename(output_file_partial, output_file)

if self.printed:
print("\n")

return 0

def urlopen(self, url, headers):
headers["Range"] = f"bytes={self.file_size}-"
request = urllib.request.Request(url, headers=headers)
try:
self.response = urllib.request.urlopen(request)
except urllib.error.HTTPError as e:
print(f"Request failed: {e.code}", file=sys.stderr)
return 1
except urllib.error.URLError as e:
print(f"Network error: {e.reason}", file=sys.stderr)
return 1

if self.response.status not in (200, 206):
print(f"Request failed: {self.response.status}", file=sys.stderr)

return 1

return 0

def perform_download(self, file, progress):
self.total_to_download += self.file_size
self.now_downloaded = 0
self.start_time = time.time()
while True:
data = self.response.read(1024)
if not data:
break

size = file.write(data)
if progress:
self.update_progress(size)

def human_readable_time(self, seconds):
hrs = int(seconds) // 3600
mins = (int(seconds) % 3600) // 60
secs = int(seconds) % 60
width = 10
if hrs > 0:
return f"{hrs}h {mins:02}m {secs:02}s".rjust(width)
elif mins > 0:
return f"{mins}m {secs:02}s".rjust(width)
else:
return f"{secs}s".rjust(width)

def human_readable_size(self, size):
width = 10
for unit in ["B", "KB", "MB", "GB", "TB"]:
if size < 1024:
return f"{size:.2f} {unit}".rjust(width)

size /= 1024

return f"{size:.2f} PB".rjust(width)

def get_terminal_width(self):
return shutil.get_terminal_size().columns

def generate_progress_prefix(self, percentage):
return f"{percentage}% |".rjust(6)

def generate_progress_suffix(self, now_downloaded_plus_file_size, speed, estimated_time):
return f"{self.human_readable_size(now_downloaded_plus_file_size)}/{self.human_readable_size(self.total_to_download)}{self.human_readable_size(speed)}/s{self.human_readable_time(estimated_time)}" # noqa: E501

def calculate_progress_bar_width(self, progress_prefix, progress_suffix):
progress_bar_width = self.get_terminal_width() - len(progress_prefix) - len(progress_suffix) - 3
if progress_bar_width < 1:
progress_bar_width = 1

return progress_bar_width

def generate_progress_bar(self, progress_bar_width, percentage):
pos = (percentage * progress_bar_width) // 100
progress_bar = ""
for i in range(progress_bar_width):
progress_bar += "█" if i < pos else " "

return progress_bar

def set_resume_point(self, output_file):
if output_file and os.path.exists(output_file):
return os.path.getsize(output_file)

return 0

def print_progress(self, progress_prefix, progress_bar, progress_suffix):
print(f"\r{progress_prefix}{progress_bar}| {progress_suffix}", end="")

def update_progress(self, chunk_size):
self.now_downloaded += chunk_size
now_downloaded_plus_file_size = self.now_downloaded + self.file_size
percentage = (now_downloaded_plus_file_size * 100) // self.total_to_download if self.total_to_download else 100
progress_prefix = self.generate_progress_prefix(percentage)
speed = self.calculate_speed(self.now_downloaded, self.start_time)
tim = (self.total_to_download - self.now_downloaded) // speed
progress_suffix = self.generate_progress_suffix(now_downloaded_plus_file_size, speed, tim)
progress_bar_width = self.calculate_progress_bar_width(progress_prefix, progress_suffix)
progress_bar = self.generate_progress_bar(progress_bar_width, percentage)
self.print_progress(progress_prefix, progress_bar, progress_suffix)
self.printed = True

def calculate_speed(self, now_downloaded, start_time):
now = time.time()
elapsed_seconds = now - start_time
return now_downloaded / elapsed_seconds

0 comments on commit f3312c9

Please sign in to comment.