Skip to content

Commit 2c7aa9f

Browse files
BeautyyuYanliwylswz
BeautyyuYanli
and
wylswz
authored
feat: auto-detect pip mirror (#750)
* feat: auto-detect pip mirror * fix: include pip mirror helper in server build * chore: clarify pip mirror auto-detect log * use latency based detection * improve documentation * make probe timeout configurable --------- Co-authored-by: yunlu.wen <yunlu.wen@dify.ai>
1 parent 9d31389 commit 2c7aa9f

8 files changed

Lines changed: 324 additions & 12 deletions

File tree

‎.env.example‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -235,5 +235,15 @@ DIFY_BACKWARDS_INVOCATION_WRITE_TIMEOUT=5000
235235
# dify backwards invocation read timeout in milliseconds
236236
DIFY_BACKWARDS_INVOCATION_READ_TIMEOUT=240000
237237
PLUGIN_IGNORE_UV_LOCK=false
238+
# auto-detect pip mirror at daemon startup by measuring latency to official PyPI and each candidate
239+
# a candidate is selected only if it is faster than official PyPI by defined amount; if no candidate meets
240+
# this threshold, no mirror is used
241+
# explicit PIP_MIRROR_URL takes precedence and disables auto-detection
242+
PIP_MIRROR_AUTO_DETECT=true
238243
PIP_MIRROR_URL=
239-
PYTHON_COMPILE_ALL_EXTRA_ARGS=
244+
# timeout in seconds for each mirror probe request (default: 5; increase on high-latency networks)
245+
PIP_MIRROR_AUTO_DETECT_TIMEOUT=5
246+
# comma-separated list of PyPI mirror URLs to probe; overrides the built-in default candidates
247+
# example: PIP_MIRROR_CANDIDATES=https://mirrors.aliyun.com/pypi/simple/,https://pypi.tuna.tsinghua.edu.cn/simple/
248+
PIP_MIRROR_CANDIDATES=
249+
PYTHON_COMPILE_ALL_EXTRA_ARGS=

‎README.md‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -128,10 +128,17 @@ Using PYTHON_COMPILE_ALL_EXTRA_ARGS="-x \.venv" prevents compileall from compili
128128

129129
```shell
130130
PLUGIN_IGNORE_UV_LOCK=true
131-
PIP_MIRROR_URL=https://mirrors.aliyun.com/pypi/simple/
131+
PIP_MIRROR_AUTO_DETECT=true
132+
PIP_MIRROR_URL=
132133
```
133134

134-
Setting PLUGIN_IGNORE_UV_LOCK=true allows uv to ignore the uv.lock file and use the configured PyPI mirror for dependency resolution.
135+
Setting PLUGIN_IGNORE_UV_LOCK=true allows uv to ignore the uv.lock file and use the selected PyPI mirror for dependency resolution.
136+
137+
PIP_MIRROR_AUTO_DETECT is enabled by default. When enabled, the daemon probes official PyPI and each candidate mirror in parallel at startup, then selects the fastest candidate. If official PyPI responds faster than all candidates, no mirror is used.
138+
139+
PIP_MIRROR_URL manually specifies the PyPI mirror. When it is set, it takes precedence over auto-detection.
140+
141+
PIP_MIRROR_CANDIDATES is an optional comma-separated list of mirror URLs to probe. When unset, the built-in defaults (`https://mirrors.aliyun.com/pypi/simple/` and `https://pypi.tuna.tsinghua.edu.cn/simple/`) are used.
135142

136143
## Benchmark
137144

‎cmd/server/main.go‎

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@ package main
22

33
import (
44
"context"
5+
"math"
6+
"net/http"
57
"time"
68

79
"github.com/kelseyhightower/envconfig"
@@ -10,6 +12,18 @@ import (
1012
"github.com/langgenius/dify-plugin-daemon/pkg/utils/log"
1113
)
1214

15+
const (
16+
officialPypiURL = "https://pypi.org/simple/"
17+
// a candidate must be at least this much faster than official PyPI (relative)
18+
// to be selected; prevents choosing a mirror based on measurement noise.
19+
pipMirrorMinImprovementFactor = 0.10
20+
)
21+
22+
var defaultMirrorCandidates = []string{
23+
"https://mirrors.aliyun.com/pypi/simple/",
24+
"https://pypi.tuna.tsinghua.edu.cn/simple/",
25+
}
26+
1327
func main() {
1428
var config app.Config
1529

@@ -39,6 +53,8 @@ func main() {
3953
}
4054
defer log.RecoverAndExit()
4155

56+
applyPipMirrorAutoDetect(&config)
57+
4258
if err = config.Validate(); err != nil {
4359
log.Panic("invalid configuration", "error", err)
4460
}
@@ -55,3 +71,91 @@ func main() {
5571

5672
(&server.App{}).Run(&config)
5773
}
74+
75+
func applyPipMirrorAutoDetect(config *app.Config) {
76+
candidates := config.PipMirrorCandidates
77+
if len(candidates) == 0 {
78+
candidates = defaultMirrorCandidates
79+
}
80+
81+
timeout := time.Duration(config.PipMirrorAutoDetectTimeoutSeconds) * time.Second
82+
mirror := detectAndApplyPipMirror(config, &http.Client{Timeout: timeout}, candidates, officialPypiURL, timeout)
83+
if mirror != "" {
84+
log.Info(
85+
"IMPORTANT: pip mirror auto-detect selected a mirror; set PIP_MIRROR_AUTO_DETECT=false to disable or PIP_MIRROR_URL=<mirror_url> to override",
86+
"mirror_url", mirror,
87+
)
88+
}
89+
}
90+
91+
func detectAndApplyPipMirror(config *app.Config, client *http.Client, candidates []string, officialURL string, timeout time.Duration) string {
92+
if !config.PipMirrorAutoDetect || config.PipMirrorUrl != "" {
93+
return ""
94+
}
95+
96+
ctx, cancel := context.WithTimeout(context.Background(), timeout)
97+
defer cancel()
98+
99+
mirror := selectFastestMirror(ctx, client, candidates, officialURL)
100+
if mirror != "" {
101+
config.PipMirrorUrl = mirror
102+
}
103+
return mirror
104+
}
105+
106+
type mirrorProbeResult struct {
107+
url string
108+
latency time.Duration
109+
ok bool
110+
}
111+
112+
func probeURL(ctx context.Context, client *http.Client, url string) mirrorProbeResult {
113+
start := time.Now()
114+
req, err := http.NewRequestWithContext(ctx, http.MethodHead, url, nil)
115+
if err != nil {
116+
return mirrorProbeResult{url: url}
117+
}
118+
resp, err := client.Do(req)
119+
if err != nil {
120+
return mirrorProbeResult{url: url}
121+
}
122+
resp.Body.Close()
123+
if resp.StatusCode >= 400 {
124+
return mirrorProbeResult{url: url}
125+
}
126+
return mirrorProbeResult{url: url, latency: time.Since(start), ok: true}
127+
}
128+
129+
func selectFastestMirror(ctx context.Context, client *http.Client, candidates []string, officialURL string) string {
130+
all := append([]string{officialURL}, candidates...)
131+
ch := make(chan mirrorProbeResult, len(all))
132+
133+
for _, u := range all {
134+
go func(u string) {
135+
ch <- probeURL(ctx, client, u)
136+
}(u)
137+
}
138+
139+
officialLatency := time.Duration(math.MaxInt64)
140+
bestURL := ""
141+
bestLatency := time.Duration(math.MaxInt64)
142+
143+
for range all {
144+
r := <-ch
145+
if !r.ok {
146+
continue
147+
}
148+
if r.url == officialURL {
149+
officialLatency = r.latency
150+
} else if r.latency < bestLatency {
151+
bestLatency = r.latency
152+
bestURL = r.url
153+
}
154+
}
155+
156+
threshold := time.Duration(float64(officialLatency) * (1 - pipMirrorMinImprovementFactor))
157+
if bestURL != "" && bestLatency < threshold {
158+
return bestURL
159+
}
160+
return ""
161+
}

‎cmd/server/pip_mirror_test.go‎

Lines changed: 138 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,138 @@
1+
package main
2+
3+
import (
4+
"context"
5+
"net/http"
6+
"net/http/httptest"
7+
"testing"
8+
"time"
9+
10+
"github.com/langgenius/dify-plugin-daemon/internal/types/app"
11+
"github.com/stretchr/testify/assert"
12+
)
13+
14+
func newInstantServer(statusCode int) *httptest.Server {
15+
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
16+
w.WriteHeader(statusCode)
17+
}))
18+
}
19+
20+
func newSlowServer(delay time.Duration, statusCode int) *httptest.Server {
21+
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
22+
time.Sleep(delay)
23+
w.WriteHeader(statusCode)
24+
}))
25+
}
26+
27+
func TestDetectAndApplyPipMirrorDisabled(t *testing.T) {
28+
config := app.Config{PipMirrorAutoDetect: false}
29+
mirror := detectAndApplyPipMirror(&config, &http.Client{Timeout: time.Second}, nil, "http://unused", time.Second)
30+
assert.Empty(t, mirror)
31+
assert.Empty(t, config.PipMirrorUrl)
32+
}
33+
34+
func TestDetectAndApplyPipMirrorExplicitURLPreserved(t *testing.T) {
35+
config := app.Config{
36+
PipMirrorAutoDetect: true,
37+
PipMirrorUrl: "https://mirror.example/simple",
38+
}
39+
mirror := detectAndApplyPipMirror(&config, &http.Client{Timeout: time.Second}, nil, "http://unused", time.Second)
40+
assert.Empty(t, mirror)
41+
assert.Equal(t, "https://mirror.example/simple", config.PipMirrorUrl)
42+
}
43+
44+
func TestDetectAndApplyPipMirrorCandidateFasterThanOfficial(t *testing.T) {
45+
official := newSlowServer(80*time.Millisecond, http.StatusOK)
46+
defer official.Close()
47+
candidate := newInstantServer(http.StatusOK)
48+
defer candidate.Close()
49+
50+
config := app.Config{PipMirrorAutoDetect: true}
51+
mirror := detectAndApplyPipMirror(&config, &http.Client{Timeout: 2 * time.Second}, []string{candidate.URL}, official.URL, 2*time.Second)
52+
assert.Equal(t, candidate.URL, mirror)
53+
assert.Equal(t, candidate.URL, config.PipMirrorUrl)
54+
}
55+
56+
func TestDetectAndApplyPipMirrorOfficialFasterNoCandidateSelected(t *testing.T) {
57+
official := newInstantServer(http.StatusOK)
58+
defer official.Close()
59+
candidate := newSlowServer(80*time.Millisecond, http.StatusOK)
60+
defer candidate.Close()
61+
62+
config := app.Config{PipMirrorAutoDetect: true}
63+
mirror := detectAndApplyPipMirror(&config, &http.Client{Timeout: 2 * time.Second}, []string{candidate.URL}, official.URL, 2*time.Second)
64+
assert.Empty(t, mirror)
65+
assert.Empty(t, config.PipMirrorUrl)
66+
}
67+
68+
func TestDetectAndApplyPipMirrorOfficialUnreachableCandidateSelected(t *testing.T) {
69+
candidate := newInstantServer(http.StatusOK)
70+
defer candidate.Close()
71+
72+
config := app.Config{PipMirrorAutoDetect: true}
73+
mirror := detectAndApplyPipMirror(
74+
&config,
75+
&http.Client{Timeout: 200 * time.Millisecond},
76+
[]string{candidate.URL},
77+
"http://127.0.0.1:1",
78+
200*time.Millisecond,
79+
)
80+
assert.Equal(t, candidate.URL, mirror)
81+
assert.Equal(t, candidate.URL, config.PipMirrorUrl)
82+
}
83+
84+
func TestDetectAndApplyPipMirrorPicksFastestCandidate(t *testing.T) {
85+
official := newSlowServer(120*time.Millisecond, http.StatusOK)
86+
defer official.Close()
87+
slow := newSlowServer(80*time.Millisecond, http.StatusOK)
88+
defer slow.Close()
89+
fast := newInstantServer(http.StatusOK)
90+
defer fast.Close()
91+
92+
config := app.Config{PipMirrorAutoDetect: true}
93+
mirror := detectAndApplyPipMirror(&config, &http.Client{Timeout: 2 * time.Second}, []string{slow.URL, fast.URL}, official.URL, 2*time.Second)
94+
assert.Equal(t, fast.URL, mirror)
95+
}
96+
97+
func TestDetectAndApplyPipMirrorAllUnreachableNoMirrorSet(t *testing.T) {
98+
config := app.Config{PipMirrorAutoDetect: true}
99+
mirror := detectAndApplyPipMirror(
100+
&config,
101+
&http.Client{Timeout: 200 * time.Millisecond},
102+
[]string{"http://127.0.0.1:1"},
103+
"http://127.0.0.1:2",
104+
200*time.Millisecond,
105+
)
106+
assert.Empty(t, mirror)
107+
assert.Empty(t, config.PipMirrorUrl)
108+
}
109+
110+
func TestDetectAndApplyPipMirrorCandidateErrorStatusIgnored(t *testing.T) {
111+
official := newInstantServer(http.StatusOK)
112+
defer official.Close()
113+
candidate := newInstantServer(http.StatusServiceUnavailable)
114+
defer candidate.Close()
115+
116+
config := app.Config{PipMirrorAutoDetect: true}
117+
mirror := detectAndApplyPipMirror(&config, &http.Client{Timeout: time.Second}, []string{candidate.URL}, official.URL, time.Second)
118+
assert.Empty(t, mirror)
119+
assert.Empty(t, config.PipMirrorUrl)
120+
}
121+
122+
func TestSelectFastestMirrorEmptyCandidates(t *testing.T) {
123+
official := newInstantServer(http.StatusOK)
124+
defer official.Close()
125+
126+
ctx := context.Background()
127+
result := selectFastestMirror(ctx, &http.Client{Timeout: time.Second}, nil, official.URL)
128+
assert.Empty(t, result)
129+
}
130+
131+
func TestApplyPipMirrorAutoDetectDoesNotBlockStartupOnAllFailures(t *testing.T) {
132+
config := app.Config{
133+
PipMirrorAutoDetect: true,
134+
PipMirrorCandidates: []string{"http://127.0.0.1:1"},
135+
}
136+
applyPipMirrorAutoDetect(&config)
137+
assert.Empty(t, config.PipMirrorUrl)
138+
}

‎internal/core/local_runtime/dependency_installation_test.go‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,24 @@ func TestDetectDependencyFileType(t *testing.T) {
7979
require.Equal(t, pyprojectTomlFile, fileType)
8080
})
8181

82+
t.Run("both files exist with mirror - requirements.txt takes priority", func(t *testing.T) {
83+
runtime := createTestRuntime(t, "plugin-with-both")
84+
runtime.appConfig.PipMirrorUrl = "https://mirrors.aliyun.com/pypi/simple/"
85+
86+
fileType, err := runtime.detectDependencyFileType()
87+
require.NoError(t, err)
88+
require.Equal(t, requirementsTxtFile, fileType)
89+
})
90+
91+
t.Run("only pyproject.toml with mirror - falls back to pyproject.toml", func(t *testing.T) {
92+
runtime := createTestRuntime(t, "plugin-with-pyproject")
93+
runtime.appConfig.PipMirrorUrl = "https://mirrors.aliyun.com/pypi/simple/"
94+
95+
fileType, err := runtime.detectDependencyFileType()
96+
require.NoError(t, err)
97+
require.Equal(t, pyprojectTomlFile, fileType)
98+
})
99+
82100
t.Run("neither file exists", func(t *testing.T) {
83101
runtime := createTestRuntime(t, "plugin-without-dependencies")
84102

‎internal/core/local_runtime/setup_python_environment.go‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -138,6 +138,14 @@ func (p *LocalPluginRuntime) detectDependencyFileType() (PythonDependencyFileTyp
138138
pyprojectPath := path.Join(p.State.WorkingPath, string(pyprojectTomlFile))
139139
requirementsPath := path.Join(p.State.WorkingPath, string(requirementsTxtFile))
140140

141+
// When a mirror is configured, prefer requirements.txt to avoid uv lock url
142+
// mismatches against a non-official index.
143+
if p.appConfig.PipMirrorUrl != "" {
144+
if _, err := os.Stat(requirementsPath); err == nil {
145+
return requirementsTxtFile, nil
146+
}
147+
}
148+
141149
if _, err := os.Stat(pyprojectPath); err == nil {
142150
return pyprojectTomlFile, nil
143151
}

‎internal/types/app/config.go‎

Lines changed: 12 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -201,15 +201,18 @@ type Config struct {
201201
MaxBundlePackageSize int64 `envconfig:"MAX_BUNDLE_PACKAGE_SIZE" validate:"required"`
202202
MaxServerlessTransactionTimeout int `envconfig:"MAX_SERVERLESS_TRANSACTION_TIMEOUT"`
203203

204-
PythonInterpreterPath string `envconfig:"PYTHON_INTERPRETER_PATH"`
205-
UvPath string `envconfig:"UV_PATH" default:""`
206-
PythonEnvInitTimeout int `envconfig:"PYTHON_ENV_INIT_TIMEOUT" validate:"required"`
207-
PythonCompileAllExtraArgs string `envconfig:"PYTHON_COMPILE_ALL_EXTRA_ARGS"`
208-
PipMirrorUrl string `envconfig:"PIP_MIRROR_URL"`
209-
PipPreferBinary bool `envconfig:"PIP_PREFER_BINARY" default:"true"`
210-
PipVerbose bool `envconfig:"PIP_VERBOSE" default:"true"`
211-
PipExtraArgs string `envconfig:"PIP_EXTRA_ARGS"`
212-
PluginIgnoreUvLock bool `envconfig:"PLUGIN_IGNORE_UV_LOCK" default:"false"`
204+
PythonInterpreterPath string `envconfig:"PYTHON_INTERPRETER_PATH"`
205+
UvPath string `envconfig:"UV_PATH" default:""`
206+
PythonEnvInitTimeout int `envconfig:"PYTHON_ENV_INIT_TIMEOUT" validate:"required"`
207+
PythonCompileAllExtraArgs string `envconfig:"PYTHON_COMPILE_ALL_EXTRA_ARGS"`
208+
PipMirrorAutoDetect bool `envconfig:"PIP_MIRROR_AUTO_DETECT" default:"true"`
209+
PipMirrorUrl string `envconfig:"PIP_MIRROR_URL"`
210+
PipMirrorCandidates []string `envconfig:"PIP_MIRROR_CANDIDATES"`
211+
PipMirrorAutoDetectTimeoutSeconds int `envconfig:"PIP_MIRROR_AUTO_DETECT_TIMEOUT" default:"5"`
212+
PipPreferBinary bool `envconfig:"PIP_PREFER_BINARY" default:"true"`
213+
PipVerbose bool `envconfig:"PIP_VERBOSE" default:"true"`
214+
PipExtraArgs string `envconfig:"PIP_EXTRA_ARGS"`
215+
PluginIgnoreUvLock bool `envconfig:"PLUGIN_IGNORE_UV_LOCK" default:"false"`
213216

214217
// Runtime buffer configuration (applies to both local and serverless runtimes)
215218
// These are the new generic names that should be used going forward

0 commit comments

Comments
 (0)