Fixing token within the docker image for the launcher.

This commit is contained in:
Nicolas Patry 2024-07-23 14:03:01 +00:00
parent dc05d7ba23
commit fa470bc851
No known key found for this signature in database
GPG Key ID: B154A218C20EBBCA
2 changed files with 15 additions and 6 deletions

8
Cargo.lock generated
View File

@ -3762,7 +3762,7 @@ dependencies = [
[[package]] [[package]]
name = "text-generation-benchmark" name = "text-generation-benchmark"
version = "2.1.2-dev0" version = "2.2.1-dev0"
dependencies = [ dependencies = [
"average", "average",
"clap", "clap",
@ -3783,7 +3783,7 @@ dependencies = [
[[package]] [[package]]
name = "text-generation-client" name = "text-generation-client"
version = "2.1.2-dev0" version = "2.2.1-dev0"
dependencies = [ dependencies = [
"async-trait", "async-trait",
"base64 0.22.1", "base64 0.22.1",
@ -3801,7 +3801,7 @@ dependencies = [
[[package]] [[package]]
name = "text-generation-launcher" name = "text-generation-launcher"
version = "2.1.2-dev0" version = "2.2.1-dev0"
dependencies = [ dependencies = [
"clap", "clap",
"ctrlc", "ctrlc",
@ -3820,7 +3820,7 @@ dependencies = [
[[package]] [[package]]
name = "text-generation-router" name = "text-generation-router"
version = "2.1.2-dev0" version = "2.2.1-dev0"
dependencies = [ dependencies = [
"async-stream", "async-stream",
"axum 0.7.5", "axum 0.7.5",

View File

@ -1,5 +1,8 @@
use clap::{Parser, ValueEnum}; use clap::{Parser, ValueEnum};
use hf_hub::{api::sync::Api, Repo, RepoType}; use hf_hub::{
api::sync::{Api, ApiBuilder},
Repo, RepoType,
};
use nix::sys::signal::{self, Signal}; use nix::sys::signal::{self, Signal};
use nix::unistd::Pid; use nix::unistd::Pid;
use serde::Deserialize; use serde::Deserialize;
@ -1401,7 +1404,13 @@ fn main() -> Result<(), LauncherError> {
let mut path = std::path::Path::new(&args.model_id).to_path_buf(); let mut path = std::path::Path::new(&args.model_id).to_path_buf();
let filename = if !path.exists() { let filename = if !path.exists() {
// Assume it's a hub id // Assume it's a hub id
let api = Api::new()?;
let api = if let Ok(token) = std::env::var("HF_TOKEN") {
// env variable has precedence over on file token.
ApiBuilder::new().with_token(Some(token)).build()?
} else {
Api::new()?
};
let repo = if let Some(ref revision) = args.revision { let repo = if let Some(ref revision) = args.revision {
api.repo(Repo::with_revision( api.repo(Repo::with_revision(
model_id, model_id,