Skip to content

Commit

Permalink
Fix logger in hf secret source
Browse files Browse the repository at this point in the history
  • Loading branch information
discordianfish committed Jan 24, 2025
1 parent 15b941e commit 8185a5e
Show file tree
Hide file tree
Showing 4 changed files with 12 additions and 4 deletions.
2 changes: 1 addition & 1 deletion pkg/cmd/agent/submit.go
Original file line number Diff line number Diff line change
Expand Up @@ -35,13 +35,13 @@ func NewSubmitCmd(logger *log.Logger) *cobra.Command {
name = ""
version = ""
)
submissionConfig.RegisterCredentialsProviders()

c, err := diambra.NewConfig(logger)
if err != nil {
level.Error(logger).Log("msg", err.Error())
os.Exit(1)
}
submissionConfig.RegisterCredentialsProviders(logger, c.Home)

cmd := &cobra.Command{
Use: "submit [flags] (directory | --submission.manifest=submission-manifest.yaml | docker-image) [args/command(s) ...]",
Expand Down
2 changes: 1 addition & 1 deletion pkg/cmd/agent/test.go
Original file line number Diff line number Diff line change
Expand Up @@ -24,12 +24,12 @@ const (

func NewTestCmd(logger *log.Logger) *cobra.Command {
submissionConfig := diambra.SubmissionConfig{}
submissionConfig.RegisterCredentialsProviders()
c, err := diambra.NewConfig(logger)
if err != nil {
level.Error(logger).Log("msg", err.Error())
os.Exit(1)
}
submissionConfig.RegisterCredentialsProviders(logger, c.Home)

cmd := &cobra.Command{
Use: "test [flags] {--submission.manifest submission-manifest.yaml | docker-image} [args/command(s) ...]",
Expand Down
5 changes: 3 additions & 2 deletions pkg/diambra/config.go
Original file line number Diff line number Diff line change
Expand Up @@ -254,9 +254,9 @@ func (c *SubmissionConfig) RegisterCredentialsProvider(name string, provider sec
}
c.credentialsProvider[name] = provider
}
func (c *SubmissionConfig) RegisterCredentialsProviders() {
func (c *SubmissionConfig) RegisterCredentialsProviders(logger log.Logger, home string) {
c.RegisterCredentialsProvider("git", &secretsources.GitCredentials{})
c.RegisterCredentialsProvider("huggingface", &secretsources.HuggingfaceCredentials{})
c.RegisterCredentialsProvider("huggingface", secretsources.NewHuggingfaceCredentials(logger, home))
}

func (c *SubmissionConfig) AddFlags(flags *pflag.FlagSet) {
Expand Down Expand Up @@ -346,6 +346,7 @@ func (c *SubmissionConfig) Submission(config *EnvConfig, args []string) (*client
if !ok {
return nil, fmt.Errorf("invalid value for --submission.secrets-from: %s", c.SecretsFrom)
}

switch c.SecretsFrom {
case "git":
secrets, err := secretsources.CredentialsFill(ss, manifest.Sources)
Expand Down
7 changes: 7 additions & 0 deletions pkg/secretsources/huggingface.go
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,13 @@ type HuggingfaceCredentials struct {
Home string
}

func NewHuggingfaceCredentials(logger log.Logger, home string) *HuggingfaceCredentials {
return &HuggingfaceCredentials{
logger: logger,
Home: home,
}
}

func (c *HuggingfaceCredentials) Credentials(url string) (map[string]string, error) {
cmd := exec.Command(pyarena.FindPython(), "-c", GetHuggingfaceToken)
stdout := &bytes.Buffer{}
Expand Down

0 comments on commit 8185a5e

Please sign in to comment.