feat(controller): add header for client favor oss domain (#3048)

client/starwhale/base/cloud.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import os
 import typing as t
 from copy import deepcopy
 from http import HTTPStatus
@@ -22,6 +23,7 @@
     SW_API_VERSION,
     DEFAULT_PAGE_IDX,
     DEFAULT_PAGE_SIZE,
+    ENV_SW_CLIENT_FAVORED_OSS_DOMAIN_ALIAS,
 )
 from starwhale.utils.fs import ensure_dir
 from starwhale.utils.http import ignore_error, wrap_sw_error_resp
@@ -166,6 +168,11 @@ def do_http_request(
         if token:
             _headers["Authorization"] = token
 
+        if os.environ.get(ENV_SW_CLIENT_FAVORED_OSS_DOMAIN_ALIAS):
+            _headers[ENV_SW_CLIENT_FAVORED_OSS_DOMAIN_ALIAS] = os.environ[
+                ENV_SW_CLIENT_FAVORED_OSS_DOMAIN_ALIAS
+            ]
+
         if not disable_default_content_type:
             _headers["Content-Type"] = "application/json"
 

client/starwhale/consts/__init__.py

@@ -224,3 +224,5 @@ class FileNode:
 D_ALIGNMENT_SIZE = 128  # for page cache
 
 DEFAULT_RESOURCE_POOL = "default"
+
+ENV_SW_CLIENT_FAVORED_OSS_DOMAIN_ALIAS = "SW_CLIENT_FAVORED_OSS_DOMAIN_ALIAS"

client/starwhale/core/model/copy.py

@@ -18,7 +18,11 @@
 
 from starwhale.base import cloud_blob_cache
 from starwhale.utils import console, convert_to_bytes
-from starwhale.consts import SW_API_VERSION, DEFAULT_MANIFEST_NAME
+from starwhale.consts import (
+    SW_API_VERSION,
+    DEFAULT_MANIFEST_NAME,
+    ENV_SW_CLIENT_FAVORED_OSS_DOMAIN_ALIAS,
+)
 from starwhale.proto_gen import model_package_storage_pb2 as pb2
 from starwhale.base.blob.store import LocalFileStore
 from starwhale.base.uri.project import Instance
@@ -131,6 +135,10 @@ async def _http_request(
         assert path is not None
         url = f"{instance.url}/api/{SW_API_VERSION}/{path.lstrip('/')}"
         headers["Authorization"] = instance.token
+        if os.environ.get(ENV_SW_CLIENT_FAVORED_OSS_DOMAIN_ALIAS):
+            headers[ENV_SW_CLIENT_FAVORED_OSS_DOMAIN_ALIAS] = os.environ[
+                ENV_SW_CLIENT_FAVORED_OSS_DOMAIN_ALIAS
+            ]
 
     url_iter = cloud_blob_cache.replace_url(url, replace)
     try:

server/controller/src/main/java/ai/starwhale/mlops/domain/job/split/JobSpliteratorImpl.java

@@ -19,6 +19,7 @@
 import ai.starwhale.mlops.common.util.BatchOperateHelper;
 import ai.starwhale.mlops.domain.job.JobDao;
 import ai.starwhale.mlops.domain.job.bo.Job;
+import ai.starwhale.mlops.domain.job.spec.Env;
 import ai.starwhale.mlops.domain.job.spec.JobSpecParser;
 import ai.starwhale.mlops.domain.job.spec.StepSpec;
 import ai.starwhale.mlops.domain.job.status.JobStatus;
@@ -44,8 +45,10 @@
 import java.util.stream.Collectors;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.commons.collections4.CollectionUtils;
+import org.springframework.beans.factory.annotation.Value;
 import org.springframework.stereotype.Service;
 import org.springframework.transaction.annotation.Transactional;
+import org.springframework.util.StringUtils;
 
 /**
  * split job by swds index
@@ -64,6 +67,8 @@ public class JobSpliteratorImpl implements JobSpliterator {
     private final StepMapper stepMapper;
     private final JobSpecParser jobSpecParser;
 
+    @Value("${sw.scheduler.clientFavorOssName}") String clientFavorOssName;
+
     public JobSpliteratorImpl(
             StoragePathCoordinator storagePathCoordinator,
             TaskMapper taskMapper,
@@ -110,7 +115,14 @@ public List<StepEntity> split(Job job) {
         }
         for (StepSpec stepSpec : stepSpecs) {
             boolean firstStep = CollectionUtils.isEmpty(stepSpec.getNeeds());
-
+            if (StringUtils.hasText(clientFavorOssName)) {
+                List<Env> env = stepSpec.getEnv();
+                if (null == env) {
+                    env = new LinkedList<>();
+                    stepSpec.setEnv(env);
+                }
+                env.add(new Env("SW_CLIENT_FAVORED_OSS_DOMAIN_ALIAS", clientFavorOssName));
+            }
             StepEntity stepEntity = StepEntity.builder()
                     .uuid(UUID.randomUUID().toString())
                     .jobId(job.getId())

server/controller/src/main/resources/application.yaml

@@ -75,6 +75,7 @@ sw:
   scheduler:
     impl: ${SW_SCHEDULER:k8s}
     backOffLimit: ${SW_TASK_BACKOFF_LIMIT:3}
+    clientFavorOssName: ${SW_CLIENT_FAVORED_OSS_DOMAIN_ALIAS:}
     docker:
       network: ${SW_DOCKER_CONTAINER_NETWORK:host}
       node-ip: ${SW_DOCKER_CONTAINER_NODE_IP:127.0.0.1}