fix: bohrium authorization expires

Signed-off-by: zjgemi <[email protected]>

Author: zjgemi

src/dflow/__init__.py

@@ -4,8 +4,7 @@
 from .argo_objects import ArgoStep, ArgoWorkflow
 from .code_gen import gen_code
 from .common import (CustomArtifact, HTTPArtifact, LineageClient,
-                     LocalArtifact, S3Artifact, import_func)
-from .common import jsonpickle
+                     LocalArtifact, S3Artifact, import_func, jsonpickle)
 from .config import config, s3_config, set_config, set_s3_config
 from .context import Context
 from .dag import DAG

src/dflow/code_gen.py

@@ -2,11 +2,10 @@
 import json
 
 from .common import (input_artifact_pattern, input_parameter_pattern,
-                     step_output_artifact_pattern,
+                     jsonpickle, step_output_artifact_pattern,
                      step_output_parameter_pattern,
                      task_output_artifact_pattern,
                      task_output_parameter_pattern)
-from .common import jsonpickle
 from .config import config
 from .dag import DAG
 from .io import InputArtifact, InputParameter, OutputArtifact, OutputParameter

src/dflow/config.py

@@ -1,5 +1,5 @@
-import os
 import json
+import os
 import tempfile
 
 

src/dflow/context_syntax.py

@@ -1,5 +1,5 @@
 import warnings
-from typing import Union, List, Any
+from typing import Any, List, Union
 
 dflow = Any
 

src/dflow/io.py

@@ -10,9 +10,9 @@
 from .utils import randstr, s3_config, upload_s3
 
 try:
-    from argo.workflows.client import (V1alpha1ArchiveStrategy, V1alpha1Inputs,
-                                       V1alpha1HTTPArtifact, V1alpha1Outputs,
-                                       V1alpha1RawArtifact)
+    from argo.workflows.client import (V1alpha1ArchiveStrategy,
+                                       V1alpha1HTTPArtifact, V1alpha1Inputs,
+                                       V1alpha1Outputs, V1alpha1RawArtifact)
 
     from .client import V1alpha1Artifact, V1alpha1Parameter, V1alpha1ValueFrom
 except Exception:

src/dflow/plugins/bohrium.py

@@ -50,10 +50,9 @@ def login(username=None, phone=None, password=None, bohrium_url=None):
         password = config["password"]
     if bohrium_url is None:
         bohrium_url = config["bohrium_url"]
-    if config["authorization"] is None:
-        config["authorization"] = _login(
-            bohrium_url + "/account/login", username, phone, password)
-        update_headers()
+    config["authorization"] = _login(
+        bohrium_url + "/account/login", username, phone, password)
+    update_headers()
     return config["authorization"]
 
 
@@ -290,7 +289,7 @@ def __getstate__(self):
     def __setstate__(self, d):
         self.__dict__.update(d)
 
-    def get_token(self):
+    def get_token(self, retry=1):
         import requests
         url = self.bohrium_url + "/brm/v1/storage/token"
         headers = {
@@ -309,6 +308,10 @@ def get_token(self):
             headers["Authorization"] = "Bearer " + self.authorization
         rsp = requests.get(url, headers=headers, params=params)
         if not rsp.text:
+            if retry > 0:
+                self.authorization = None
+                self.get_token(retry=retry-1)
+                return
             raise RuntimeError("Bohrium unauthorized")
         res = json.loads(rsp.text)
         _raise_error(res, "get storage token")

src/dflow/plugins/dispatcher.py

@@ -1,10 +1,10 @@
 import abc
 import json
 import os
+import re
 from abc import ABC
 from copy import deepcopy
 from getpass import getpass
-import re
 from typing import Dict, List, Optional, Union
 
 from ..common import CustomArtifact, S3Artifact

src/dflow/plugins/ray.py

@@ -1,14 +1,14 @@
 from copy import deepcopy
-from typing import Union, List, Optional, Any
+from typing import Any, List, Optional, Union
 
 from dflow.executor import Executor
 from dflow.op_template import ScriptOPTemplate
 from dflow.utils import randstr
 
 try:
-    from argo.workflows.client import (V1Volume, V1VolumeMount,
-                                       V1alpha1UserContainer,
-                                       V1EmptyDirVolumeSource)
+    from argo.workflows.client import (V1alpha1UserContainer,
+                                       V1EmptyDirVolumeSource, V1Volume,
+                                       V1VolumeMount)
 except ImportError:
     pass
 

src/dflow/slurm.py

@@ -1,7 +1,7 @@
 import os
 import re
 from copy import deepcopy
-from typing import Optional, Dict, List, Union
+from typing import Dict, List, Optional, Union
 
 from .config import config
 from .executor import (Executor, RemoteExecutor, render_script_with_tmp_root,