|
13 | 13 | from experiment_utils.providers.base import InfrastructureProvider |
14 | 14 |
|
15 | 15 |
|
| 16 | +# Experiment-config data_type -> the exporter's --data-type value. |
| 17 | +ALIBABA_DATA_TYPE_CLI_VALUES = {"node": "node", "msresource": "ms-resource"} |
| 18 | + |
| 19 | + |
16 | 20 | class ClusterDataExporterService(BaseService): |
17 | 21 | """ |
18 | 22 | Service for managing cluster_data_exporter via Docker. |
@@ -71,10 +75,11 @@ def start( |
71 | 75 | # Get number of nodes from provider (assuming it has this info) |
72 | 76 | num_nodes = kwargs.get("num_nodes", 1) |
73 | 77 |
|
74 | | - # Assert that we have exactly 2 nodes |
75 | | - assert num_nodes == 1, ( |
76 | | - f"cluster_data_exporter requires exactly 1 node (num_nodes==1), " |
77 | | - f"got {num_nodes}" |
| 78 | + # One worker node next to the coordinator; local mode reports 0 nodes |
| 79 | + # and runs everything on this machine. |
| 80 | + assert num_nodes in (0, 1), ( |
| 81 | + f"cluster_data_exporter requires one worker node (num_nodes==1) " |
| 82 | + f"or local mode (num_nodes==0), got {num_nodes}" |
78 | 83 | ) |
79 | 84 |
|
80 | 85 | # Extract configuration |
@@ -258,7 +263,8 @@ def _build_docker_command( |
258 | 263 |
|
259 | 264 | elif provider == "alibaba": |
260 | 265 | if "data_type" in config: |
261 | | - cmd_parts.append(f"--data-type={config['data_type']}") |
| 266 | + data_type = ALIBABA_DATA_TYPE_CLI_VALUES[config["data_type"]] |
| 267 | + cmd_parts.append(f"--data-type={data_type}") |
262 | 268 | if "data_year" in config: |
263 | 269 | cmd_parts.append(f"--data-year={config['data_year']}") |
264 | 270 |
|
@@ -400,23 +406,18 @@ def _validate_alibaba_data( |
400 | 406 | Raises: |
401 | 407 | ValueError: If required files are missing |
402 | 408 | """ |
403 | | - # Determine expected file pattern based on data type and year |
404 | | - if data_type == "node": |
405 | | - if data_year == 2021 or data_year == 2022: |
406 | | - pattern = "Node_*.csv.gz" |
407 | | - else: |
408 | | - raise ValueError( |
409 | | - f"Invalid data_year for Alibaba node data: {data_year}" |
410 | | - ) |
411 | | - elif data_type == "msresource": |
412 | | - if data_year == 2021 or data_year == 2022: |
413 | | - pattern = "MsResource_*.csv.gz" |
414 | | - else: |
415 | | - raise ValueError( |
416 | | - f"Invalid data_year for Alibaba msresource data: {data_year}" |
417 | | - ) |
418 | | - else: |
419 | | - raise ValueError(f"Invalid data_type for Alibaba: {data_type}") |
| 409 | + # File names the exporter reads (alibaba_metrics/{node,ms_resource}.rs). |
| 410 | + patterns = { |
| 411 | + ("node", 2021): "Node_*.csv.gz", |
| 412 | + ("node", 2022): "NodeMetrics_*.csv.gz", |
| 413 | + ("msresource", 2021): "MSResource_*.csv.gz", |
| 414 | + ("msresource", 2022): "MSMetrics_*.csv.gz", |
| 415 | + } |
| 416 | + pattern = patterns.get((data_type, int(data_year))) |
| 417 | + if pattern is None: |
| 418 | + raise ValueError( |
| 419 | + f"Invalid Alibaba data_type/data_year: {data_type}/{data_year}" |
| 420 | + ) |
420 | 421 |
|
421 | 422 | # Check for data files on remote node |
422 | 423 | target_node = self.node_offset + 1 |
@@ -497,12 +498,15 @@ def _wait_for_health( |
497 | 498 | while time.time() - start_time < timeout: |
498 | 499 | # Run curl from the remote node to check health |
499 | 500 | check_cmd = f"curl -s -o /dev/null -w '%{{http_code}}' {url}" |
| 501 | + # curl exits non-zero while the exporter is still starting; the |
| 502 | + # local provider raises on that, so keep polling instead. |
500 | 503 | result = self.provider.execute_command( |
501 | 504 | node_idx=node_idx, |
502 | 505 | cmd=check_cmd, |
503 | 506 | cmd_dir="", |
504 | 507 | nohup=False, |
505 | 508 | popen=False, |
| 509 | + ignore_errors=True, |
506 | 510 | ) |
507 | 511 |
|
508 | 512 | try: |
|
0 commit comments