Skip to content

Commit

Permalink
chore: shortcut client when result is available with errors
Browse files Browse the repository at this point in the history
Signed-off-by: mudler <mudler@localai.io>
  • Loading branch information
mudler committed Dec 16, 2024
1 parent d0d5e01 commit 27ea9ae
Show file tree
Hide file tree
Showing 3 changed files with 30 additions and 26 deletions.
33 changes: 18 additions & 15 deletions cmd/tee-worker/config.go
Original file line number Diff line number Diff line change
Expand Up @@ -11,40 +11,43 @@ import (
)

func readConfig() types.JobConfiguration {
// The jobs will then unmarshal from this configuration to the specific configuration
// that is needed for the job
jc := types.JobConfiguration{}

dataDir := os.Getenv("DATA_DIR")
if dataDir == "" {
dataDir = "/home/masa"
os.Setenv("DATA_DIR", dataDir)
}

jc["data_dir"] = dataDir

// Read the env file
if err := godotenv.Load(filepath.Join(dataDir, ".env")); err != nil {
fmt.Println("Failed reading env file!")
panic(err)
}

webScraperBlacklist := os.Getenv("WEBSCRAPER_BLACKLIST")

blacklistURLs := strings.Split(webScraperBlacklist, ",")
for i, u := range blacklistURLs {
blacklistURLs[i] = strings.TrimSpace(u)
if webScraperBlacklist != "" {
blacklistURLs := strings.Split(webScraperBlacklist, ",")
for i, u := range blacklistURLs {
blacklistURLs[i] = strings.TrimSpace(u)
}
jc["webscraper_blacklist"] = blacklistURLs
}

twitterAccount := os.Getenv("TWITTER_ACCOUNTS")
if twitterAccount != "" {
twitterAccounts := strings.Split(twitterAccount, ",")
for i, u := range twitterAccounts {
twitterAccounts[i] = strings.TrimSpace(u)
}

twitterAccounts := strings.Split(twitterAccount, ",")
for i, u := range twitterAccounts {
twitterAccounts[i] = strings.TrimSpace(u)
jc["twitter_accounts"] = twitterAccounts
}

// Read the .env file and set the global configuration for all the jobs
// The jobs will then unmarshal from this configuration to the specific configuration
// that is needed for the job
jc := types.JobConfiguration{}
jc["webscraper_blacklist"] = blacklistURLs
jc["twitter_accounts"] = twitterAccounts
jc["data_dir"] = dataDir

return jc
}

Expand Down
16 changes: 8 additions & 8 deletions pkg/client/http.go
Original file line number Diff line number Diff line change
Expand Up @@ -86,27 +86,27 @@ func (c *Client) Decrypt(encryptedResult string) (string, error) {
}

// GetJobResult retrieves the encrypted result of a job.
func (c *Client) GetResult(jobUUID string) (string, error) {
func (c *Client) GetResult(jobUUID string) (string, bool, error) {
resp, err := c.HTTPClient.Get(c.BaseURL + "/job/" + jobUUID)
if err != nil {
return "", fmt.Errorf("error sending GET request: %w", err)
return "", false, fmt.Errorf("error sending GET request: %w", err)
}
defer resp.Body.Close()

body, err := io.ReadAll(resp.Body)
if err != nil {
return "", fmt.Errorf("error reading response body: %w", err)
return "", false, fmt.Errorf("error reading response body: %w", err)
}

if resp.StatusCode == http.StatusNotFound {
err = fmt.Errorf("job not found or not ready")
return "", false, fmt.Errorf("job not found")
}

if resp.StatusCode != http.StatusOK {
respErr := types.JobError{}
json.Unmarshal(body, &respErr)
respErr := types.JobError{}
json.Unmarshal(body, &respErr)
if respErr.Error != "" {
err = fmt.Errorf("error: %s", respErr.Error)
}

return string(body), err
return string(body), true, err
}
7 changes: 4 additions & 3 deletions pkg/client/result.go
Original file line number Diff line number Diff line change
Expand Up @@ -21,22 +21,23 @@ func (jr *JobResult) SetDelay(delay time.Duration) {
}

// GetJobResult retrieves the encrypted result of a job.
func (jr *JobResult) getResult() (string, error) {
func (jr *JobResult) getResult() (string, bool, error) {
return jr.client.GetResult(jr.UUID)
}

// Get polls the server until the job result is ready or a timeout occurs.
func (jr *JobResult) Get() (result string, err error) {
retries := 0
var resultIsAvailable bool

for {
if retries >= jr.maxRetries {
return "", fmt.Errorf("max retries reached: %w", err)
}
retries++

result, err = jr.getResult()
if err == nil {
result, resultIsAvailable, err = jr.getResult()
if err == nil || resultIsAvailable {
break
}
time.Sleep(jr.delay)
Expand Down

0 comments on commit 27ea9ae

Please sign in to comment.