在 Google Colab 中使用 JuiceFS

JuiceFS發表於2024-03-22

Google Colaboratory(Colab)是一個由 Google 提供的雲端 Jupyter 程式設計筆記本,直接透過瀏覽器即可進行 Python 程式設計。Colab 充分利用谷歌的閒置雲端計算資源,為公眾提供免費的的線上程式設計服務,以及免費的 GPU 資源,雖然在使用方面有一定的規則限制,但對於一般的研究和學習來說綽綽有餘。

訪問 Colab,可以新建筆記本,也可以從 Google Drive、Github 載入筆記本,或直接從本地上傳。

Colab 免費提供的 Python 程式設計環境十分慷慨,如下圖,足有 12 GB 的記憶體和 100 GB 的硬碟。

如果需要使用硬體加速,可以切換執行時型別,筆者的賬戶可以免費使用 T4 GPU 和 TPU。

不過需要注意,Colab 執行時是臨時的,平臺會監測執行時的活動狀態,長時間的閒置和長時間的高強度使用,都會導致執行時被收回釋放,所有資料都會被清空。

Colab 上的資料持久化

為了讓 Colab 可以持久化地儲存檔案,人們通常會選擇使用 Google Drive。如下圖,使用時在介面左側的檔案管理中點選按鈕即可將 Google Drive 掛載到執行時,把需要長期保留或重複使用的資料儲存在裡面,再次使用可以從 Google Drive 中載入,這就避免了執行被釋放時丟失資料。

除了 Google Drive 以外,你還可以使用 JuiceFS 作為 Colab 筆記本的持久化儲存,從而更為靈活地儲存和共享更大規模的資料。

JuiceFS 與 Google Drive

這裡先提供一個表格供讀者參考,後文會展開介紹 JuiceFS 的技術架構以及如何建立一個適用於 Google Colab 的檔案系統。

簡言之,Google Drive 有平臺優勢,更容易整合到 Colab,也有多種容量規格以供擴容,但在使用上會有一些限制,比如單位時間的總上傳量,總檔案數量等。而 JuiceFS 是自建服務,沒有此類限制,而且在費用方面可以透過靈活地組織資源來降低是用成本。

JuiceFS Google Drive
價格 彈性費用(取決於後設資料引擎和物件儲存的費用) 按固定容量訂閱
整合到 Colab 簡單 簡單
是否需要維護 需要 不需要
擴容能力 無容量上限 15GB ~ 30TB
上傳限制 無限制 24 小時內可向雲端硬碟上傳和複製 750 GB 資料
跨平臺共享 靈活 一般

使用 JuiceFS

JuiceFS 是面向雲的高效能分散式檔案系統,它在 Apache-2.0 協議下開源,具有完備的 POSIX 相容性,並支援 FUSE POSIX、HDFS、S3、Kubernetes CSI Driver、WebDAV 等多種訪問方式。

在 Colab 中可以直接採用 FUSE POSIX 方式,以守護程序形式掛載到執行時中使用。

技術架構

一個典型的 JuiceFS 檔案系統由一個負責儲存資料的物件儲存和一個負責儲存後設資料的後設資料引擎組成。

在物件儲存方面,JuiceFS 支援幾乎所有公有云物件儲存、私有部署的物件儲存、NFS、SFTP 以及本地磁碟等。在後設資料引擎方面,支援 Redis、Postgres、MySQL、TiKV、SQLite 等多種資料庫。

因為採用了資料與後設資料分離儲存的架構,JuiceFS 檔案系統的讀和寫操作會先在後設資料引擎上進行查詢和處理,只有涉及到實際資料時才會訪問物件儲存,這樣就能極高效地處理海量的資料,相比於直接與物件儲存互動,JuiceFS 會有更好的效能表現。

簡言之,後設資料引擎至關重要,想要提升 JuiceFS 檔案系統的效能,一個黃金法則是“儘量靠近業務端部署後設資料引擎”。以 Colab 為例,它的伺服器大多位於美國,所以找一個美國的雲伺服器來部署 Redis 並搭配一個可靠的物件儲存是構建 JuiceFS 檔案系統相對理想的搭配。

適用於 Colab 的組合

後設資料引擎

為了在 Colab 上儘量發揮 JuiceFS 的潛能,筆者選擇了一個位於美國矽谷的雲伺服器,將它用於部署 Redis 作為 JuiceFS 的後設資料引擎,配置如下:

  • 位置:美國矽谷
  • CPU:2 核
  • RAM:4GB
  • SSD: 60GB
  • 系統:Ubuntu Server 22.04
  • IP:18.18.18.18
  • 域名:redis.xxx.com

注:上述 IP 和域名均為演示目的編寫,如有需要請替換成你的真實資訊

按照 JuiceFS 官方提供的資料,採用 Redis 這類鍵值資料庫作為後設資料引擎時,儲存佔用大概為 300 位元組/檔案,那麼 1GB 記憶體大概可以儲存 350 萬個檔案的後設資料,讀者可以根據預期的檔案總量來決定給伺服器配置多少記憶體。

在本文中,筆者使用 Docker 部署 Redis,並透過 Let's Encrypt 簽發了一個免費的 SSL 證書對服務端進行加密:

# 拉取 redis 映象
sudo docker pull bitnami/redis:7.2

# 刪除已存在的同名 redis 容器(如果存在的話)
sudo docker rm -f redis

# 建立新的 redis 容器
sudo docker run -d --name redis \
-p 6379:6379 \
-v redis_aof_data:/bitnami/redis/data \
-v ./ssl:/ssl \
-e REDIS_PASSWORD=abcdefg \
-e REDIS_TLS_ENABLED=yes \
-e REDIS_TLS_PORT_NUMBER=6379 \
-e REDIS_TLS_AUTH_CLIENTS=no \
-e REDIS_TLS_CERT_FILE=/ssl/redis.xxx.com.crt \
-e REDIS_TLS_KEY_FILE=/ssl/redis.xxx.com.key \
-e REDIS_TLS_CA_FILE=/ssl/ca.crt \
--restart unless-stopped \
bitnami/redis:7.2

該 Redis 例項啟用了以下功能:

  • AOF(追加寫到本地檔案):它會將每個寫操作都記錄到本地磁碟,從而提高資料的安全性。
  • 服務端 SSL:Redis 伺服器會使用 SSL/TLS 協議與客戶端進行通訊。將 SSL 證書放在 ssl 目錄中即可,注意修改環境變數中的證書檔名稱。

Redis 部署完畢後,還需要檢查防火牆,確保伺服器開放了 6379 埠的入站請求。這樣,Redis 後設資料引擎就準備完畢了。

物件儲存

物件儲存方面,筆者選擇使用 Cloudflare R2,因為它上下行流量均免費,只需為儲存和 API 請求付費,很適合 Colab 這種需要在外部訪問 JuiceFS 的場景。特別是儲存量較大的情況,每次都要將模型資料載入到 Colab 執行時,下行收費會產生不小的開支。

以下是演示目的編寫的物件儲存資訊,請在實際配置時替換成自己的真實資訊:

  • Bucket Name:myjfs
  • Endpoint URL:https://xxx.r2.cloudflarestorage.com/myjfs
  • Access Key:abcdefg
  • Secret Key:gfedcba

建立檔案系統

後設資料引擎和物件儲存都準備好了,接下來使用 JuiceFS 客戶端來建立檔案系統。

這個步驟可以在任何一臺支援安裝 JuiceFS 客戶端的電腦上執行,可以是你本地正在使用的電腦或是部署了 Redis 例項的那臺雲服務。

因為 JuiceFS 是基於雲的,只要 JuiceFS 客戶端能夠訪問到後設資料引擎和物件儲存就可以建立和使用。

首先,安裝 JuiceFS 客戶端:

# macOS 或 Linux 系統
curl -sSL https://d.juicefs.com/install | sh -

# Windows 系統(建議使用 Scoop)
scoop install juicefs

其他系統及安裝方法請參考 JuiceFS 官方安裝文件

然後,使用已準備的後設資料引擎和物件儲存來建立檔案系統:

# 建立檔案系統
juicefs format --storage s3 \
--bucket https://xxx.r2.cloudflarestorage.com/myjfs \
--access-key abcdefg \
--secret-key gfedcba \
rediss://:abcdefg@redis.xxx.com/1 \
myjfs

JuiceFS 檔案系統只需要建立一次,然後就可以在任何安裝了 JuiceFS 客戶端的裝置上掛載和使用,它是基於雲的共享檔案系統。可以多裝置、跨地區、跨網路同時讀寫訪問。

現在你可以在任何安裝了 JuiceFS 客戶端的裝置上掛載使用這個檔案系統,以下是幾種常用的訪問方式:

# 以 FUSE POSIX 方式掛載
juicefs mount rediss://:abcdefg@redis.xxx.com/1 mnt

# 以 S3 Gateway 形式掛載
export MINIO_ROOT_USER=admin  
export MINIO_ROOT_PASSWORD=12345678
juicefs gateway rediss://:abcdefg@redis.xxx.com/1 localhost:9000

# 以 WebDAV 形式掛載
juicefs webdav rediss://:abcdefg@redis.xxx.com/1 localhost:8000

可以看到,掛載 JuiceFS 檔案系統時只需指定後設資料引擎 URL,不再需要物件儲存相關的資訊。這是因為在建立檔案系統的時候,物件儲存相關的資訊已經被寫入了後設資料引擎。

在 Colab 中掛載 JuiceFS

如下圖,Colab 執行時的底層是一個 Ubuntu 系統,所以,只需要在 Colab 上安裝 JuiceFS 客戶端,執行掛載命令即可使用。

可以將安裝命令和掛載命令放在一個程式碼塊中,也可以像下圖這樣,將它們分成兩個獨立的程式碼塊。

請注意,掛載 JuiceFS 時不要忘記 -d 選項,它的作用是讓 JuiceFS 以守護程序的方式掛載到後臺。因為 Colab 每次只允許一個程式碼塊執行,如果不將 JuiceFS 掛載到後臺,他就會一直讓程式碼塊處於執行狀態,導致其他程式碼塊無法執行。

如下圖,左側檔案管理器中可以看到已掛載的 JuiceFS 檔案系統。

使用舉例

例一:用 JuiceFS 儲存 Fooocus 模型

Fooocus 是開源的 AI 圖片生成器,底層仍然使用 Stable Diffusion 模型,但將複雜的引數進行了調優和封裝,讓使用者可以獲得像 Midjourney 一樣簡單直觀的作圖體驗。

你可以直接使用 Fooocus 官方提供的 Colab Notebook,在其基礎上新增安裝和掛載 JuiceFS 檔案系統的程式碼塊。

也可以參考以下程式碼更靈活地建立和管理 Fooocus 相關的程式碼:

# 安裝 JuiceFS 客戶端
!curl -sSL https://d.juicefs.com/install | sh -

# 掛載 JuiceFS 檔案系統
!juicefs mount rediss://:abcdefg@redis.xxx.com/1 myjfs -d

# 在 JuiceFS 中建立 Fooocus 模型目錄結構
!mkdir -p myjfs/models/{checkpoints,loras,embeddings,vae_approx,upscale_models,inpaint,controlnet,clip_vision}

# 克隆 Fooocus 倉庫
!git clone https://github.com/lllyasviel/Fooocus.git

在 Fooocus 專案根目錄建立一個自定義的 config.txt 檔案,讓 Fooocus 以 JuiceFS 中的目錄作為預設的模型儲存目錄:

{
  "path_checkpoints": "/content/myjfs/models/checkpoints",
  "path_loras": "/content/myjfs/models/loras",
  "path_embeddings": "/content/myjfs/models/embeddings",
  "path_vae_approx": "/content/myjfs/models/vae_approx",
  "path_upscale_models": "/content/myjfs/models/upscale_models",
  "path_inpaint": "/content/myjfs/models/inpaint",
  "path_controlnet": "/content/myjfs/models/controlnet",
  "path_clip_vision": "/content/myjfs/models/clip_vision"
}

啟動 Fooocus

!pip install pygit2==1.12.2
%cd /content/Fooocus
!python entry_with_update.py --share

初次使用需要從公共倉庫下載模型,這會需要一些時間,你可以在本地同時掛載 JuiceFS 檔案系統,觀察模型儲存的情況。

再次使用時,只需掛載 JuiceFS 檔案系統並確保 Fooocus 可以從中讀取模型。程式會動態地從 JuiceFS 拉取所需的模型,儘管這仍然需要一些時間,但相比每次都從公共倉庫完全重新下載要更加方便。特別是模型經過微調或產生了自定義資料情況,使用 JuiceFS 儲存相應資料的優勢就會更加明顯。

例二:用 JuiceFS 儲存 Chroma 向量資料庫

在 Colab 上構建 RAG(檢索增強生成)應用也是比較常見的,這往往涉及到要把各種資料生成的 embedding 資料儲存到向量資料庫。

Llamaindex 預設採用 OpenAI 的 text-embedding 模型對輸入的資料進行向量化,如果不想每次都重新生成 embedding 資料,就需要將這些資料儲存到向量資料庫。比如使用開源的 Chroma 向量資料庫,因為它預設將資料儲存在本地磁碟,在 Colab 中需要注意資料庫的儲存位置,以防執行時收回造成資料丟失。

這裡筆者提供一組 Colab 筆記本程式碼,讓你可以將 Llamaindex 生成的 embedding 儲存到 Chroma 資料庫,而這個 Chroma 資料庫將完全儲存到 JuiceFS。

# 安裝 JuiceFS 客戶端
!curl -sSL https://d.juicefs.com/install | sh -

# 掛載 JuiceFS 檔案系統
!juicefs mount rediss://:abcdefg@redis.xxx.com/1 myjfs -d

# 安裝 Llamaindex 和 chroma 相關的包
!pip install llama-index chromadb kaleido python-multipart pypdf cohere

# 從 Colab 環境變數讀取 OpenAI API 金鑰
from google.colab import userdata
import openai

openai.api_key = userdata.get('OPENAI_API_KEY')

把需要轉換成 embedding 的檔案放在 myjfs/data/ 目錄中,執行以下程式碼生成 embedding 並儲存到 Chroma。

得益於 JuiceFS 基於雲的共享訪問特性,可以同時在本地掛載 JuiceFS 並將所需的資料放入相應的目錄。

import chromadb
from llama_index import VectorStoreIndex, SimpleDirectoryReader
from llama_index.vector_stores import ChromaVectorStore
from llama_index.storage.storage_context import StorageContext

# load some documents
documents = SimpleDirectoryReader("./myjfs/data").load_data()

# initialize client, setting path to save data
db = chromadb.PersistentClient(path="./myjfs/chroma_db")

# create collection
chroma_collection = db.get_or_create_collection("great_ceo")

# assign chroma as the vector_store to the context
vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
storage_context = StorageContext.from_defaults(vector_store=vector_store)

# create your index
index = VectorStoreIndex.from_documents(
    documents, storage_context=storage_context,
    embed_model_name="text-embedding-3-small",
)

使用時,直接讓 Chroma 從 JuiceFS 讀取資料。

import chromadb
from llama_index import VectorStoreIndex, SimpleDirectoryReader
from llama_index.vector_stores import ChromaVectorStore
from llama_index.storage.storage_context import StorageContext

# initialize client, setting path to save data
db = chromadb.PersistentClient(path="./myjfs/chroma_db")

# create collection
chroma_collection = db.get_or_create_collection("great_ceo")

# assign chroma as the vector_store to the context
vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
storage_context = StorageContext.from_defaults(vector_store=vector_store)

# 從向量資料庫載入索引
index = VectorStoreIndex.from_vector_store(
    vector_store, storage_context=storage_context
)

測試一下用自己 Index 作為知識庫與 GPT 進行對話:

# create a query engine and query
query_engine = index.as_query_engine()
response = query_engine.query("這本書講了什麼?")
print(response)

這樣一來,每次進入新的 Colab 執行時,只要掛載 JuiceFS 就可以直接使用這些已建立的向量資料。其實,不止是在 Colab,任何需要訪問這些向量資料的地方都可以透過掛載 JuiceFS 來使用。

總結

本文介紹瞭如何在 Google Colab 中使用 JuiceFS 來持久化儲存資料,透過例項介紹瞭如何為 JuiceFS 準備後設資料引擎和物件儲存來儘量發揮它的效能,以及在 Colab 中的安裝和掛載方法。最後透過 Fooocus 和 Chroma 兩個例子,演示了在實際應用中如何利用 JuiceFS 來更好地儲存並重複利用資料。

希望這篇文章的內容能夠對你起到一定的幫助,如果你有任何疑問,歡迎在評論區留言討論。如果你對 JuiceFS 有興趣,可以檢視官方文件瞭解更多用法和效能調優方面的內容。

希望這篇內容能夠對你有一些幫助,如果有其他疑問歡迎加入 JuiceFS 社群與大家共同交流。

相關文章